Use linear bits for decoding/encoding (#40)

src/json/decoder.rs

@@ -16,7 +16,7 @@ use tket_json_rs::circuit_json::SerialCircuit;
 
 use super::op::JsonOp;
 use super::{try_param_to_constant, METADATA_IMPLICIT_PERM, METADATA_PHASE};
-use crate::utils::{BIT, QB};
+use crate::utils::{LINEAR_BIT, QB};
 
 /// The state of an in-progress [`DFGBuilder`] being built from a [`SerialCircuit`].
 ///
@@ -42,10 +42,6 @@ impl JsonDecoder {
         let num_qubits = serialcirc.qubits.len();
         let num_bits = serialcirc.bits.len();
 
-        if num_bits > 0 {
-            unimplemented!("TKET1's linear bits are not supported yet.");
-        }
-
         // Map each (register name, index) pair to an offset in the signature.
         let mut wire_map: HashMap<RegisterHash, usize> =
             HashMap::with_capacity(num_bits + num_qubits);
@@ -61,7 +57,8 @@ impl JsonDecoder {
             }
             wire_map.insert((register, 0).into(), i);
         }
-        let sig = Signature::new_linear([vec![QB; num_qubits], vec![BIT; num_bits]].concat());
+        let sig =
+            Signature::new_linear([vec![QB; num_qubits], vec![LINEAR_BIT; num_bits]].concat());
 
         let mut dfg = DFGBuilder::new(sig.input, sig.output).unwrap();
 

src/json/encoder.rs

@@ -3,14 +3,13 @@
 use std::collections::HashMap;
 
 use hugr::ops::{Const, ConstValue, OpType};
-use hugr::types::SimpleType;
 use hugr::Wire;
-use itertools::{Either, Itertools};
+use itertools::Itertools;
 use tket_json_rs::circuit_json::{self, Permutation, SerialCircuit};
 
 use crate::circuit::command::{CircuitUnit, Command};
 use crate::circuit::Circuit;
-use crate::utils::{BIT, QB};
+use crate::utils::{LINEAR_BIT, QB};
 
 use super::op::JsonOp;
 use super::{OpConvertError, METADATA_IMPLICIT_PERM, METADATA_PHASE};
@@ -26,9 +25,10 @@ pub(super) struct JsonEncoder {
     implicit_permutation: Vec<Permutation>,
     /// The current commands
     commands: Vec<circuit_json::Command>,
-    /// The linear units of the circuit, mapped to their TKET1 registers and
-    /// types.
-    units: HashMap<CircuitUnit, (circuit_json::Register, SimpleType)>,
+    /// The TKET1 qubit registers associated to each qubit unit of the circuit.
+    qubit_regs: HashMap<CircuitUnit, circuit_json::Register>,
+    /// The TKET1 bit registers associated to each linear bit unit of the circuit.
+    bit_regs: HashMap<CircuitUnit, circuit_json::Register>,
     /// A register of wires with constant values, used to recover TKET1
     /// parameters.
     parameters: HashMap<Wire, String>,
@@ -39,33 +39,35 @@ impl JsonEncoder {
     pub fn new<'circ>(circ: &impl Circuit<'circ>) -> Self {
         let name = circ.name().map(str::to_string);
 
-        // Compute the linear qubit and bit registers
-        // TODO We are checking for Hugr's bit. We should change this when the
-        // decoding starts using linear bits instead.
-        // TODO Through an error on non-recognized unit types, or just ignore?
-        let units = circ
-            .units()
-            .into_iter()
-            .filter_map(|(u, ty)| {
-                let CircuitUnit::Linear(index) = u else { return None; };
-                if ty != QB && ty != BIT {
-                    return None;
-                }
-                let prefix = if ty == QB { "q" } else { "b" };
-                let reg = circuit_json::Register(prefix.to_string(), vec![index as i64]);
-                Some((u, (reg, ty)))
-            })
-            .collect();
+        // Compute the linear qubit and bit registers. Each one have independent
+        // indices starting from zero.
+        //
+        // TODO Throw an error on non-recognized unit types, or just ignore?
+        let mut bit_units = HashMap::new();
+        let mut qubit_units = HashMap::new();
+        for (unit, ty) in circ.units() {
+            if ty == QB {
+                let index = vec![qubit_units.len() as i64];
+                let reg = circuit_json::Register("q".to_string(), index);
+                qubit_units.insert(unit, reg);
+            } else if ty == LINEAR_BIT {
+                let index = vec![bit_units.len() as i64];
+                let reg = circuit_json::Register("c".to_string(), index);
+                bit_units.insert(unit, reg);
+            }
+        }
 
         let mut encoder = Self {
             name,
             phase: "0".to_string(),
             implicit_permutation: vec![],
             commands: vec![],
-            units,
+            qubit_regs: qubit_units,
+            bit_regs: bit_units,
             parameters: HashMap::new(),
         };
 
+        // Encode other parameters stored in the metadata
         if let Some(meta) = circ.get_metadata(circ.root()).as_object() {
             if let Some(phase) = meta.get(METADATA_PHASE) {
                 // TODO: Check for invalid encoded metadata
@@ -89,7 +91,7 @@ impl JsonEncoder {
         let args = command
             .inputs
             .iter()
-            .filter_map(|u| self.units.get(u).map(|(reg, _)| reg.clone()))
+            .filter_map(|&u| self.unit_to_register(u))
             .collect();
 
         // TODO Restore the opgroup (once the decoding supports it)
@@ -107,19 +109,12 @@ impl JsonEncoder {
     }
 
     pub fn finish(self) -> SerialCircuit {
-        let (qubits, bits) =
-            self.units
-                .into_iter()
-                .partition_map(|(_unit, (reg, ty))| match ty == QB {
-                    true => Either::Left(reg),
-                    false => Either::Right(reg),
-                });
         SerialCircuit {
             name: self.name,
             phase: self.phase,
             commands: self.commands,
-            qubits,
-            bits,
+            qubits: self.qubit_regs.into_values().collect_vec(),
+            bits: self.bit_regs.into_values().collect_vec(),
             implicit_permutation: self.implicit_permutation,
         }
     }
@@ -168,4 +163,11 @@ impl JsonEncoder {
             }
         }
     }
+
+    fn unit_to_register(&self, unit: CircuitUnit) -> Option<circuit_json::Register> {
+        self.qubit_regs
+            .get(&unit)
+            .or_else(|| self.bit_regs.get(&unit))
+            .cloned()
+    }
 }

src/json/op.rs

@@ -16,7 +16,7 @@ use tket_json_rs::optype::OpType as JsonOpType;
 
 use super::{try_param_to_constant, OpConvertError};
 use crate::resource::try_unwrap_json_op;
-use crate::utils::{BIT, F64, QB};
+use crate::utils::{F64, LINEAR_BIT, QB};
 
 /// A serialized operation, containing the operation type and all its attributes.
 ///
@@ -121,12 +121,9 @@ impl JsonOp {
     }
 
     /// Compute the signature of the operation.
-    //
-    // TODO: We are using Hugr's non-liner bits. We should have a custom linear
-    // bit type instead.
     #[inline]
     pub fn signature(&self) -> Signature {
-        let linear = [vec![QB; self.num_qubits], vec![BIT; self.num_bits]].concat();
+        let linear = [vec![QB; self.num_qubits], vec![LINEAR_BIT; self.num_bits]].concat();
         let params = vec![F64; self.num_params];
         Signature::new_df([linear.clone(), params].concat(), linear)
     }
@@ -235,7 +232,7 @@ impl TryFrom<&OpType> for JsonOp {
                 LeafOp::CX => JsonOpType::CX,
                 LeafOp::ZZMax => JsonOpType::ZZMax,
                 LeafOp::Reset => JsonOpType::Reset,
-                LeafOp::Measure => JsonOpType::Measure,
+                //LeafOp::Measure => JsonOpType::Measure,
                 LeafOp::T => JsonOpType::T,
                 LeafOp::S => JsonOpType::S,
                 LeafOp::X => JsonOpType::X,
@@ -253,8 +250,8 @@ impl TryFrom<&OpType> for JsonOp {
                 // CustomOp is handled above
                 _ => return Err(err()),
             },
-            OpType::Input(_) => JsonOpType::Input,
-            OpType::Output(_) => JsonOpType::Output,
+            //OpType::Input(_) => JsonOpType::Input,
+            //OpType::Output(_) => JsonOpType::Output,
             //hugr::ops::OpType::FuncDefn(_) => todo!(),
             //hugr::ops::OpType::FuncDecl(_) => todo!(),
             //hugr::ops::OpType::Const(_) => todo!(),
@@ -271,7 +268,7 @@ impl TryFrom<&OpType> for JsonOp {
         for ty in op.signature().input.iter() {
             if *ty == QB {
                 num_qubits += 1;
-            } else if *ty == BIT {
+            } else if *ty == LINEAR_BIT {
                 num_bits += 1;
             } else if *ty == F64 {
                 num_params += 1;

src/json/tests.rs

@@ -40,19 +40,21 @@ fn read_json_unknown_op() {
     // custom and output
 
     let circ_s = r#"{
-        "bits": [],
+        "phase": "1/2",
+        "bits": [["c", [0]], ["c", [1]]],
+        "qubits": [["q", [0]], ["q", [1]], ["q", [2]]],
         "commands": [
-            {"args": [["q", [0]], ["q", [1]], ["q", [2]]], "op": {"type": "CSWAP"}}
+            {"args": [["q", [0]], ["q", [1]], ["q", [2]]], "op": {"type": "CSWAP"}},
+            {"args": [["q", [1]], ["c", [1]]], "op": {"type": "Measure"}},
+            {"args": [["q", [2]], ["c", [0]]], "op": {"type": "Measure"}}
         ],
         "created_qubits": [],
         "discarded_qubits": [],
-        "implicit_permutation": [[["q", [0]], ["q", [0]]], [["q", [1]], ["q", [1]]], [["q", [2]], ["q", [2]]]],
-        "phase": "0",
-        "qubits": [["q", [0]], ["q", [1]], ["q", [2]]]
+        "implicit_permutation": [[["q", [0]], ["q", [0]]], [["q", [1]], ["q", [1]]], [["q", [2]], ["q", [2]]]]
     }"#;
 
     let ser: SerialCircuit = serde_json::from_str(circ_s).unwrap();
-    assert_eq!(ser.commands.len(), 1);
+    assert_eq!(ser.commands.len(), 3);
 
     let hugr: Hugr = ser.clone().decode().unwrap();
     let circ = FlatRegionView::new(&hugr, hugr.root());

src/utils.rs

@@ -1,9 +1,11 @@
 //! Utility functions for the library.
 
-use hugr::types::{ClassicType, SimpleType};
+use hugr::types::{ClassicType, CustomType, SimpleType};
+use smol_str::SmolStr;
 
 pub(crate) const QB: SimpleType = SimpleType::Qubit;
-pub(crate) const BIT: SimpleType = SimpleType::Classic(ClassicType::Int(1));
+pub(crate) const LINEAR_BIT: SimpleType =
+    SimpleType::Qpaque(CustomType::new_simple(SmolStr::new_inline("LBit")));
 pub(crate) const F64: SimpleType = SimpleType::Classic(ClassicType::F64);
 
 #[allow(dead_code)]