diff --git a/src/hugr/hugrmut.rs b/src/hugr/hugrmut.rs index 1f5dd6d30..597c8f1f6 100644 --- a/src/hugr/hugrmut.rs +++ b/src/hugr/hugrmut.rs @@ -315,13 +315,11 @@ impl + AsMut> HugrMut for T { let src_port = self .get_optype(src) .other_output_port() - .expect("Source operation has no non-dataflow outgoing edges") - .as_outgoing()?; + .expect("Source operation has no non-dataflow outgoing edges"); let dst_port = self .get_optype(dst) .other_input_port() - .expect("Destination operation has no non-dataflow incoming edges") - .as_incoming()?; + .expect("Destination operation has no non-dataflow incoming edges"); self.connect(src, src_port, dst, dst_port)?; Ok((src_port, dst_port)) } diff --git a/src/hugr/serialize.rs b/src/hugr/serialize.rs index b49549b0f..770b8c383 100644 --- a/src/hugr/serialize.rs +++ b/src/hugr/serialize.rs @@ -148,11 +148,10 @@ impl TryFrom<&Hugr> for SerHugrV0 { .expect("Could not reach one of the nodes"); let find_offset = |node: Node, offset: usize, dir: Direction, hugr: &Hugr| { - let sig = hugr.signature(node).unwrap_or_default(); - let offset = match offset < sig.port_count(dir) { - true => Some(offset as u16), - false => None, - }; + let op = hugr.get_optype(node); + let is_value_port = offset < op.value_port_count(dir); + let is_static_input = op.static_port(dir).map_or(false, |p| p.index() == offset); + let offset = (is_value_port || is_static_input).then_some(offset as u16); (node_rekey[&node], offset) }; @@ -263,6 +262,8 @@ pub mod test { use crate::hugr::NodeType; use crate::ops::custom::{ExtensionOp, OpaqueOp}; use crate::ops::{dataflow::IOTrait, Input, LeafOp, Module, Output, DFG}; + use crate::std_extensions::arithmetic::float_ops::FLOAT_OPS_REGISTRY; + use crate::std_extensions::arithmetic::float_types::{ConstF64, FLOAT64_TYPE}; use crate::std_extensions::logic::NotOp; use crate::types::{FunctionType, Type}; use crate::OutgoingPort; @@ -480,4 +481,18 @@ pub mod test { new_hugr.validate(&EMPTY_REG).unwrap_err(); new_hugr.validate(&PRELUDE_REGISTRY).unwrap(); } + + #[test] + fn constants_roundtrip() -> Result<(), Box> { + let mut builder = DFGBuilder::new(FunctionType::new(vec![], vec![FLOAT64_TYPE])).unwrap(); + let w = builder.add_load_const(ConstF64::new(0.5))?; + let hugr = builder.finish_hugr_with_outputs([w], &FLOAT_OPS_REGISTRY)?; + + let ser = serde_json::to_string(&hugr)?; + let deser = serde_json::from_str(&ser)?; + + assert_eq!(hugr, deser); + + Ok(()) + } } diff --git a/src/ops.rs b/src/ops.rs index 9c44f7f00..b2b373791 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -10,7 +10,7 @@ pub mod module; pub mod tag; pub mod validate; use crate::extension::ExtensionSet; -use crate::types::{EdgeKind, FunctionType, Type}; +use crate::types::{EdgeKind, FunctionType}; use crate::{Direction, OutgoingPort, Port}; use crate::{IncomingPort, PortIndex}; use paste::paste; @@ -109,11 +109,12 @@ impl Default for OpType { } impl OpType { - /// The edge kind for the non-dataflow or constant ports of the - /// operation, not described by the signature. + /// The edge kind for the non-dataflow ports of the operation, not described + /// by the signature. /// - /// If not None, a single extra multiport of that kind will be present on - /// the given direction. + /// If not None, a single extra port of that kind will be present on + /// the given direction after any dataflow or constant ports. + #[inline] pub fn other_port_kind(&self, dir: Direction) -> Option { match dir { Direction::Incoming => self.other_input(), @@ -121,21 +122,46 @@ impl OpType { } } + /// The edge kind for the static ports of the operation, not described by + /// the dataflow signature. + /// + /// If not None, an extra input port of that kind will be present on the + /// given direction after any dataflow ports and before any + /// [`OpType::other_port_kind`] ports. + #[inline] + pub fn static_port_kind(&self, dir: Direction) -> Option { + match dir { + Direction::Incoming => self.static_input(), + Direction::Outgoing => self.static_output(), + } + } + /// Returns the edge kind for the given port. + /// + /// The result may be a value port, a static port, or a non-dataflow port. + /// See [`OpType::dataflow_signature`], [`OpType::static_port_kind`], and + /// [`OpType::other_port_kind`]. pub fn port_kind(&self, port: impl Into) -> Option { let signature = self.dataflow_signature().unwrap_or_default(); let port: Port = port.into(); - let port_as_in = port.as_incoming().ok(); let dir = port.direction(); - let port_count = signature.port_count(dir); + + // Dataflow ports if port.index() < port_count { - signature.port_type(port).cloned().map(EdgeKind::Value) - } else if port_as_in.is_some() && port_as_in == self.static_input_port() { - Some(EdgeKind::Static(static_in_type(self))) - } else { - self.other_port_kind(dir) + return signature.port_type(port).cloned().map(EdgeKind::Value); } + + // Constant port + let static_kind = self.static_port_kind(dir); + if port.index() == port_count { + if let Some(kind) = static_kind { + return Some(kind); + } + } + + // Non-dataflow ports + self.other_port_kind(dir) } /// The non-dataflow port for the operation, not described by the signature. @@ -155,7 +181,47 @@ impl OpType { } } + /// The non-dataflow input port for the operation, not described by the signature. + /// See `[OpType::other_port]`. + #[inline] + pub fn other_input_port(&self) -> Option { + self.other_port(Direction::Incoming) + .map(|p| p.as_incoming().unwrap()) + } + + /// The non-dataflow output port for the operation, not described by the signature. + /// See `[OpType::other_port]`. + #[inline] + pub fn other_output_port(&self) -> Option { + self.other_port(Direction::Outgoing) + .map(|p| p.as_outgoing().unwrap()) + } + + /// If the op has a static port, the port of that input. + /// + /// See [`OpType::static_input_port`] and [`OpType::static_output_port`]. + #[inline] + pub fn static_port(&self, dir: Direction) -> Option { + self.static_port_kind(dir)?; + Some(Port::new(dir, self.value_port_count(dir))) + } + + /// If the op has a static input ([`Call`] and [`LoadConstant`]), the port of that input. + #[inline] + pub fn static_input_port(&self) -> Option { + self.static_port(Direction::Incoming) + .map(|p| p.as_incoming().unwrap()) + } + + /// If the op has a static output ([`Const`], [`FuncDefn`], [`FuncDecl`]), the port of that output. + #[inline] + pub fn static_output_port(&self) -> Option { + self.static_port(Direction::Outgoing) + .map(|p| p.as_outgoing().unwrap()) + } + /// The number of Value ports in given direction. + #[inline] pub fn value_port_count(&self, dir: portgraph::Direction) -> usize { self.dataflow_signature() .map(|sig| sig.port_count(dir)) @@ -163,76 +229,44 @@ impl OpType { } /// The number of Value input ports. + #[inline] pub fn value_input_count(&self) -> usize { self.value_port_count(Direction::Incoming) } /// The number of Value output ports. + #[inline] pub fn value_output_count(&self) -> usize { self.value_port_count(Direction::Outgoing) } - /// The non-dataflow input port for the operation, not described by the signature. - /// See `[OpType::other_port]`. - pub fn other_input_port(&self) -> Option { - self.other_port(Direction::Incoming) - } - - /// The non-dataflow input port for the operation, not described by the signature. - /// See `[OpType::other_port]`. - pub fn other_output_port(&self) -> Option { - self.other_port(Direction::Outgoing) - } - - /// If the op has a static input (Call and LoadConstant), the port of that input. - pub fn static_input_port(&self) -> Option { - match self { - OpType::Call(call) => Some(call.called_function_port()), - OpType::LoadConstant(l) => Some(l.constant_port()), - _ => None, - } - } - - /// If the op has a static output (Const, FuncDefn, FuncDecl), the port of that output. - pub fn static_output_port(&self) -> Option { - OpTag::StaticOutput - .is_superset(self.tag()) - .then_some(0.into()) - } - /// Returns the number of ports for the given direction. + #[inline] pub fn port_count(&self, dir: Direction) -> usize { + let has_static_port = self.static_port_kind(dir).is_some(); let non_df_count = self.non_df_port_count(dir); - // if there is a static input it comes before the non_df_ports - let static_input = - (dir == Direction::Incoming && OpTag::StaticInput.is_superset(self.tag())) as usize; - self.value_port_count(dir) + non_df_count + static_input + self.value_port_count(dir) + has_static_port as usize + non_df_count } /// Returns the number of inputs ports for the operation. + #[inline] pub fn input_count(&self) -> usize { self.port_count(Direction::Incoming) } /// Returns the number of outputs ports for the operation. + #[inline] pub fn output_count(&self) -> usize { self.port_count(Direction::Outgoing) } /// Checks whether the operation can contain children nodes. + #[inline] pub fn is_container(&self) -> bool { self.validity_flags().allowed_children != OpTag::None } } -fn static_in_type(op: &OpType) -> Type { - match op { - OpType::Call(call) => Type::new_function(call.called_function_type().clone()), - OpType::LoadConstant(load) => load.constant_type().clone(), - _ => panic!("this function should not be called if the optype is not known to be Call or LoadConst.") - } -} - /// Macro used by operations that want their /// name to be the same as their type name macro_rules! impl_op_name { @@ -286,10 +320,10 @@ pub trait OpTrait { ExtensionSet::new() } - /// The edge kind for the non-dataflow or constant inputs of the operation, + /// The edge kind for the non-dataflow inputs of the operation, /// not described by the signature. /// - /// If not None, a single extra output multiport of that kind will be + /// If not None, a single extra input port of that kind will be /// present. fn other_input(&self) -> Option { None @@ -298,12 +332,30 @@ pub trait OpTrait { /// The edge kind for the non-dataflow outputs of the operation, not /// described by the signature. /// - /// If not None, a single extra output multiport of that kind will be + /// If not None, a single extra output port of that kind will be /// present. fn other_output(&self) -> Option { None } + /// The edge kind for a single constant input of the operation, not + /// described by the dataflow signature. + /// + /// If not None, an extra input port of that kind will be present after the + /// dataflow input ports and before any [`OpTrait::other_input`] ports. + fn static_input(&self) -> Option { + None + } + + /// The edge kind for a single constant output of the operation, not + /// described by the dataflow signature. + /// + /// If not None, an extra output port of that kind will be present after the + /// dataflow input ports and before any [`OpTrait::other_output`] ports. + fn static_output(&self) -> Option { + None + } + /// Get the number of non-dataflow multiports. fn non_df_port_count(&self, dir: Direction) -> usize { match dir { diff --git a/src/ops/constant.rs b/src/ops/constant.rs index c66ab6f76..64954a655 100644 --- a/src/ops/constant.rs +++ b/src/ops/constant.rs @@ -105,7 +105,7 @@ impl OpTrait for Const { ::TAG } - fn other_output(&self) -> Option { + fn static_output(&self) -> Option { Some(EdgeKind::Static(self.typ.clone())) } } diff --git a/src/ops/dataflow.rs b/src/ops/dataflow.rs index d830fd0d2..d9d2bcd71 100644 --- a/src/ops/dataflow.rs +++ b/src/ops/dataflow.rs @@ -17,6 +17,7 @@ pub(crate) trait DataflowOpTrait { /// /// If not None, a single extra output multiport of that kind will be /// present. + #[inline] fn other_input(&self) -> Option { Some(EdgeKind::StateOrder) } @@ -25,9 +26,20 @@ pub(crate) trait DataflowOpTrait { /// /// If not None, a single extra output multiport of that kind will be /// present. + #[inline] fn other_output(&self) -> Option { Some(EdgeKind::StateOrder) } + + /// The edge kind for a single constant input of the operation, not + /// described by the dataflow signature. + /// + /// If not None, an extra input port of that kind will be present after the + /// dataflow input ports and before any [`DataflowOpTrait::other_input`] ports. + #[inline] + fn static_input(&self) -> Option { + None + } } /// Helpers to construct input and output nodes @@ -125,6 +137,10 @@ impl OpTrait for T { fn other_output(&self) -> Option { DataflowOpTrait::other_output(self) } + + fn static_input(&self) -> Option { + DataflowOpTrait::static_input(self) + } } impl StaticTag for T { const TAG: OpTag = T::TAG; @@ -132,9 +148,9 @@ impl StaticTag for T { /// Call a function directly. /// -/// The first ports correspond to the signature of the function being -/// called. Immediately following those ports, the first input port is -/// connected to the def/declare block with a `ConstE` edge. +/// The first ports correspond to the signature of the function being called. +/// The port immediately following those those is connected to the def/declare +/// block with a [`EdgeKind::Static`] edge. #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct Call { /// Signature of function being called @@ -152,6 +168,11 @@ impl DataflowOpTrait for Call { fn signature(&self) -> FunctionType { self.signature.clone() } + + fn static_input(&self) -> Option { + let fn_typ = Type::new_function(self.called_function_type().clone()); + Some(EdgeKind::Static(fn_typ)) + } } impl Call { #[inline] @@ -161,6 +182,21 @@ impl Call { } /// The IncomingPort which links to the function being called. + /// + /// This matches [`OpType::static_input_port`]. + /// + /// ``` + /// # use hugr::ops::dataflow::Call; + /// # use hugr::ops::OpType; + /// # use hugr::types::FunctionType; + /// # use hugr::extension::prelude::QB_T; + /// let signature = FunctionType::new(vec![QB_T, QB_T], vec![QB_T, QB_T]); + /// let call = Call { signature }; + /// let op = OpType::Call(call.clone()); + /// assert_eq!(op.static_input_port(), Some(call.called_function_port())); + /// ``` + /// + /// [`OpType::static_input_port`]: crate::ops::OpType::static_input_port #[inline] pub fn called_function_port(&self) -> IncomingPort { self.called_function_type().input_count().into() @@ -208,6 +244,10 @@ impl DataflowOpTrait for LoadConstant { fn signature(&self) -> FunctionType { FunctionType::new(TypeRow::new(), vec![self.datatype.clone()]) } + + fn static_input(&self) -> Option { + Some(EdgeKind::Static(self.constant_type().clone())) + } } impl LoadConstant { #[inline] @@ -217,6 +257,20 @@ impl LoadConstant { } /// The IncomingPort which links to the loaded constant. + /// + /// This matches [`OpType::static_input_port`]. + /// + /// ``` + /// # use hugr::ops::dataflow::LoadConstant; + /// # use hugr::ops::OpType; + /// # use hugr::types::Type; + /// let datatype = Type::UNIT; + /// let load_constant = LoadConstant { datatype }; + /// let op = OpType::LoadConstant(load_constant.clone()); + /// assert_eq!(op.static_input_port(), Some(load_constant.constant_port())); + /// ``` + /// + /// [`OpType::static_input_port`]: crate::ops::OpType::static_input_port #[inline] pub fn constant_port(&self) -> IncomingPort { 0.into() diff --git a/src/ops/module.rs b/src/ops/module.rs index 571789183..41ca76c44 100644 --- a/src/ops/module.rs +++ b/src/ops/module.rs @@ -52,7 +52,7 @@ impl OpTrait for FuncDefn { ::TAG } - fn other_output(&self) -> Option { + fn static_output(&self) -> Option { Some(EdgeKind::Static(Type::new_function(self.signature.clone()))) } } @@ -79,7 +79,7 @@ impl OpTrait for FuncDecl { ::TAG } - fn other_output(&self) -> Option { + fn static_output(&self) -> Option { Some(EdgeKind::Static(Type::new_function(self.signature.clone()))) } }