tket_qsystem/extension/
random.rs

1//! This module defines the "tket.qsystem.random" extension, which includes
2//! random number generation (RNG) functions available for Quantinuum systems.
3
4use std::sync::{Arc, Weak};
5
6use derive_more::derive::Display;
7use hugr::{
8    Extension, Wire,
9    builder::{BuildError, Dataflow},
10    extension::{
11        ExtensionBuildError, ExtensionId, ExtensionRegistry, OpDef, PRELUDE, SignatureFunc,
12        TypeDefBound, Version,
13        prelude::{UnwrapBuilder, option_type},
14        simple_op::{MakeOpDef, MakeRegisteredOp, try_from_name},
15    },
16    std_extensions::arithmetic::{float_types::float64_type, int_types::int_type},
17    types::{CustomType, Signature, Type, TypeBound},
18};
19use lazy_static::lazy_static;
20use smol_str::SmolStr;
21use strum::{EnumIter, EnumString, IntoStaticStr};
22
23/// The extension ID for the RNG extension.
24pub const EXTENSION_ID: ExtensionId = ExtensionId::new_unchecked("tket.qsystem.random");
25/// The version of the "tket.qsystem.random" extension.
26pub const EXTENSION_VERSION: Version = Version::new(0, 2, 1);
27
28lazy_static! {
29    /// The "tket.qsystem.random" extension.
30    pub static ref EXTENSION: Arc<Extension> = {
31         Extension::new_arc(EXTENSION_ID, EXTENSION_VERSION, |ext, ext_ref| {
32            add_random_type_defs(ext, ext_ref).unwrap();
33            RandomOp::load_all_ops( ext, ext_ref).unwrap();
34        })
35    };
36
37    /// Extension registry including the "tket.qsystem.random" extension and
38    /// dependencies.
39    pub static ref REGISTRY: ExtensionRegistry = ExtensionRegistry::new([
40        EXTENSION.to_owned(),
41        PRELUDE.to_owned(),
42    ]);
43
44    /// The name of the `tket.qsystem.random.context` type.
45    pub static ref CONTEXT_TYPE_NAME: SmolStr = SmolStr::new_inline("context");
46}
47
48fn add_random_type_defs(
49    extension: &mut Extension,
50    extension_ref: &Weak<Extension>,
51) -> Result<(), ExtensionBuildError> {
52    extension.add_type(
53        CONTEXT_TYPE_NAME.to_owned(),
54        vec![],
55        "The linear RNG context type".into(),
56        TypeDefBound::any(),
57        extension_ref,
58    )?;
59    Ok(())
60}
61
62/// An enum for RNG extension types.
63pub enum RandomType {
64    /// The linear RNG context type.
65    RNGContext,
66}
67
68impl RandomType {
69    fn get_type(&self, extension_ref: &Weak<Extension>) -> Type {
70        match self {
71            Self::RNGContext { .. } => CustomType::new(
72                CONTEXT_TYPE_NAME.to_owned(),
73                vec![],
74                EXTENSION_ID,
75                EXTENSION_VERSION,
76                TypeBound::Linear,
77                extension_ref,
78            ),
79        }
80        .into()
81    }
82}
83
84#[derive(
85    Clone,
86    Copy,
87    Debug,
88    serde::Serialize,
89    serde::Deserialize,
90    Hash,
91    PartialEq,
92    Eq,
93    PartialOrd,
94    Ord,
95    EnumIter,
96    IntoStaticStr,
97    EnumString,
98    Display,
99)]
100#[non_exhaustive]
101/// The operations provided by the random extension.
102pub enum RandomOp {
103    /// `fn random_int(RNGContext) -> (u32, RNGContext)`
104    RandomInt,
105    /// `fn random_float(RNGContext) -> (f32, RNGContext)`
106    RandomFloat,
107    /// `fn random_int_bounded(RNGContext, bound: u32) -> (u32, RNGContext)`
108    RandomIntBounded,
109    /// `fn new_rng_context(seed: u64) -> Option<RNGContext>` // return None on second call
110    NewRNGContext,
111    /// `fn delete_rng_context(RNGContext) -> ()`
112    DeleteRNGContext,
113    /// `fn random_advance(RNGContext, delta: u64) -> RNGContext`
114    RandomAdvance,
115}
116
117impl MakeOpDef for RandomOp {
118    fn opdef_id(&self) -> hugr::ops::OpName {
119        <&'static str>::from(self).into()
120    }
121
122    fn init_signature(&self, extension_ref: &std::sync::Weak<Extension>) -> SignatureFunc {
123        match self {
124            RandomOp::RandomInt => Signature::new(
125                vec![RandomType::RNGContext.get_type(extension_ref)],
126                vec![int_type(5), RandomType::RNGContext.get_type(extension_ref)],
127            ),
128            RandomOp::RandomFloat => Signature::new(
129                vec![RandomType::RNGContext.get_type(extension_ref)],
130                vec![
131                    float64_type(),
132                    RandomType::RNGContext.get_type(extension_ref),
133                ],
134            ),
135            RandomOp::RandomIntBounded => Signature::new(
136                vec![RandomType::RNGContext.get_type(extension_ref), int_type(5)],
137                vec![int_type(5), RandomType::RNGContext.get_type(extension_ref)],
138            ),
139            RandomOp::RandomAdvance => Signature::new(
140                vec![RandomType::RNGContext.get_type(extension_ref), int_type(6)],
141                vec![RandomType::RNGContext.get_type(extension_ref)],
142            ),
143            RandomOp::NewRNGContext => Signature::new(
144                vec![int_type(6)],
145                vec![Type::from(option_type(vec![
146                    RandomType::RNGContext.get_type(extension_ref),
147                ]))],
148            ),
149            RandomOp::DeleteRNGContext => {
150                Signature::new(vec![RandomType::RNGContext.get_type(extension_ref)], vec![])
151            }
152        }
153        .into()
154    }
155
156    fn from_def(op_def: &OpDef) -> Result<Self, hugr::extension::simple_op::OpLoadError> {
157        try_from_name(op_def.name(), op_def.extension_id())
158    }
159
160    fn extension(&self) -> ExtensionId {
161        EXTENSION_ID
162    }
163
164    fn extension_ref(&self) -> std::sync::Weak<Extension> {
165        Arc::downgrade(&EXTENSION)
166    }
167
168    fn description(&self) -> String {
169        match self {
170            RandomOp::RandomInt => "Generate a random 32-bit unsigned integer.",
171            RandomOp::RandomFloat => "Generate a random floating point value in the range [0,1).",
172            RandomOp::RandomIntBounded => "Generate a random 32-bit unsigned integer less than `bound`.",
173            RandomOp::RandomAdvance => "Advance or backtrack the RNG state by `delta` steps",
174            RandomOp::NewRNGContext => {
175                "Seed the RNG and return a new RNG context. Required before using other RNG ops, can be called only once."
176            }
177            RandomOp::DeleteRNGContext => "Discard the given RNG context.",
178        }
179        .to_string()
180    }
181}
182
183impl MakeRegisteredOp for RandomOp {
184    fn extension_id(&self) -> ExtensionId {
185        EXTENSION_ID
186    }
187
188    fn extension_ref(&self) -> Arc<Extension> {
189        EXTENSION.clone()
190    }
191}
192
193/// An extension trait for [Dataflow] providing methods to add
194/// "tket.qsystem.random" operations.
195pub trait RandomOpBuilder: Dataflow + UnwrapBuilder {
196    /// Add a "tket.qsystem.random.random_int" op.
197    fn add_random_int(&mut self, ctx: Wire) -> Result<[Wire; 2], BuildError> {
198        Ok(self
199            .add_dataflow_op(RandomOp::RandomInt, [ctx])?
200            .outputs_arr())
201    }
202
203    /// Add a "tket.qsystem.random.random_float" op.
204    fn add_random_float(&mut self, ctx: Wire) -> Result<[Wire; 2], BuildError> {
205        Ok(self
206            .add_dataflow_op(RandomOp::RandomFloat, [ctx])?
207            .outputs_arr())
208    }
209
210    /// Add a "tket.qsystem.random.random_int_bounded" op.
211    fn add_random_int_bounded(&mut self, ctx: Wire, bound: Wire) -> Result<[Wire; 2], BuildError> {
212        Ok(self
213            .add_dataflow_op(RandomOp::RandomIntBounded, [ctx, bound])?
214            .outputs_arr())
215    }
216
217    /// Add a "tket.qsystem.random.random_advance" op.
218    fn add_random_advance(&mut self, ctx: Wire, delta: Wire) -> Result<Wire, BuildError> {
219        Ok(self
220            .add_dataflow_op(RandomOp::RandomAdvance, [ctx, delta])?
221            .out_wire(0))
222    }
223
224    /// Add a "tket.qsystem.random.new_rng_context" op.
225    fn add_new_rng_context(&mut self, seed: Wire) -> Result<Wire, BuildError> {
226        Ok(self
227            .add_dataflow_op(RandomOp::NewRNGContext, [seed])?
228            .out_wire(0))
229    }
230
231    /// Add a "tket.qsystem.random.delete_rng_context" op.
232    fn add_delete_rng_context(&mut self, ctx: Wire) -> Result<(), BuildError> {
233        self.add_dataflow_op(RandomOp::DeleteRNGContext, [ctx])?;
234        Ok(())
235    }
236}
237
238impl<D: Dataflow> RandomOpBuilder for D {}
239
240#[cfg(test)]
241mod test {
242    use hugr::extension::simple_op::MakeExtensionOp;
243    use hugr::ops::Value;
244    use hugr::std_extensions::arithmetic::int_types::ConstInt;
245
246    use hugr::HugrView;
247    use hugr::builder::{DataflowHugr, FunctionBuilder};
248    use strum::IntoEnumIterator;
249
250    use super::*;
251
252    #[test]
253    fn create_extension() {
254        assert_eq!(EXTENSION.name(), &EXTENSION_ID);
255
256        for o in RandomOp::iter() {
257            assert_eq!(
258                RandomOp::from_def(EXTENSION.get_op(&o.op_id()).unwrap()),
259                Ok(o)
260            );
261        }
262    }
263
264    #[test]
265    fn test_random_op_builder() {
266        let hugr = {
267            let mut func_builder = FunctionBuilder::new(
268                "random_op_builder",
269                Signature::new(vec![], vec![int_type(5)]),
270            )
271            .unwrap();
272
273            let seed =
274                func_builder.add_load_const(Value::from(ConstInt::new_u(6, 123456).unwrap()));
275            let maybe_ctx = func_builder.add_new_rng_context(seed).unwrap();
276            let [ctx] = func_builder
277                .build_unwrap_sum(
278                    1,
279                    option_type(vec![
280                        RandomType::RNGContext.get_type(&Arc::downgrade(&EXTENSION)),
281                    ]),
282                    maybe_ctx,
283                )
284                .unwrap();
285            let bound = func_builder.add_load_const(Value::from(ConstInt::new_u(5, 100).unwrap()));
286            let delta = func_builder.add_load_const(Value::from(ConstInt::new_s(6, -1).unwrap()));
287            let [_, ctx] = func_builder.add_random_int_bounded(ctx, bound).unwrap();
288            let [_, ctx] = func_builder.add_random_float(ctx).unwrap();
289            let ctx = func_builder.add_random_advance(ctx, delta).unwrap();
290            let [rnd, ctx] = func_builder.add_random_int(ctx).unwrap();
291            func_builder.add_delete_rng_context(ctx).unwrap();
292            func_builder.finish_hugr_with_outputs([rnd]).unwrap()
293        };
294        hugr.validate().unwrap()
295    }
296}
tket_qsystem/extension/random.rs

tket_qsystem/extension/
random.rs
