feat: Introduce new enum for constant folding; deprecate Value::Function

hugr-core/src/builder/build_traits.rs

@@ -119,7 +119,7 @@ pub trait Container {
     }
 
     /// Insert a copy of a HUGR as a child of the container.
-    fn add_hugr_view(&mut self, child: &impl HugrView) -> InsertionResult {
+    fn add_hugr_view<H: HugrView>(&mut self, child: &H) -> InsertionResult<H::Node, Node> {
         let parent = self.container_node();
         self.hugr_mut().insert_from_view(parent, child)
     }

hugr-core/src/extension.rs

@@ -19,9 +19,8 @@ use derive_more::Display;
 use thiserror::Error;
 
 use crate::hugr::IdentList;
-use crate::ops::constant::{ValueName, ValueNameRef};
 use crate::ops::custom::{ExtensionOp, OpaqueOp};
-use crate::ops::{self, OpName, OpNameRef};
+use crate::ops::{OpName, OpNameRef};
 use crate::types::type_param::{TypeArg, TypeArgError, TypeParam};
 use crate::types::RowVariable;
 use crate::types::{check_typevar_decl, CustomType, Substitution, TypeBound, TypeName};
@@ -34,7 +33,7 @@ pub mod resolution;
 pub mod simple_op;
 mod type_def;
 
-pub use const_fold::{fold_out_row, ConstFold, ConstFoldResult, Folder};
+pub use const_fold::{fold_out_row, ConstFold, ConstFoldResult, FoldVal, Folder};
 pub use op_def::{
     CustomSignatureFunc, CustomValidator, LowerFunc, OpDef, SignatureFromArgs, SignatureFunc,
     ValidateJustArgs, ValidateTypeArgs,
@@ -378,6 +377,7 @@ pub static EMPTY_REG: ExtensionRegistry = ExtensionRegistry {
 /// TODO: decide on failure modes
 #[derive(Debug, Clone, Error, PartialEq, Eq)]
 #[allow(missing_docs)]
+#[non_exhaustive]
 pub enum SignatureError {
     /// Name mismatch
     #[error("Definition name ({0}) and instantiation name ({1}) do not match.")]
@@ -496,37 +496,6 @@ impl CustomConcrete for CustomType {
     }
 }
 
-/// A constant value provided by a extension.
-/// Must be an instance of a type available to the extension.
-#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
-pub struct ExtensionValue {
-    extension: ExtensionId,
-    name: ValueName,
-    typed_value: ops::Value,
-}
-
-impl ExtensionValue {
-    /// Returns a reference to the typed value of this [`ExtensionValue`].
-    pub fn typed_value(&self) -> &ops::Value {
-        &self.typed_value
-    }
-
-    /// Returns a mutable reference to the typed value of this [`ExtensionValue`].
-    pub(super) fn typed_value_mut(&mut self) -> &mut ops::Value {
-        &mut self.typed_value
-    }
-
-    /// Returns a reference to the name of this [`ExtensionValue`].
-    pub fn name(&self) -> &str {
-        self.name.as_str()
-    }
-
-    /// Returns a reference to the extension this [`ExtensionValue`] belongs to.
-    pub fn extension(&self) -> &ExtensionId {
-        &self.extension
-    }
-}
-
 /// A unique identifier for a extension.
 ///
 /// The actual [`Extension`] is stored externally.
@@ -582,8 +551,6 @@ pub struct Extension {
     pub runtime_reqs: ExtensionSet,
     /// Types defined by this extension.
     types: BTreeMap<TypeName, TypeDef>,
-    /// Static values defined by this extension.
-    values: BTreeMap<ValueName, ExtensionValue>,
     /// Operation declarations with serializable definitions.
     // Note: serde will serialize this because we configure with `features=["rc"]`.
     // That will clone anything that has multiple references, but each
@@ -607,7 +574,6 @@ impl Extension {
             version,
             runtime_reqs: Default::default(),
             types: Default::default(),
-            values: Default::default(),
             operations: Default::default(),
         }
     }
@@ -679,11 +645,6 @@ impl Extension {
         self.types.get(type_name)
     }
 
-    /// Allows read-only access to the values in this Extension
-    pub fn get_value(&self, value_name: &ValueNameRef) -> Option<&ExtensionValue> {
-        self.values.get(value_name)
-    }
-
     /// Returns the name of the extension.
     pub fn name(&self) -> &ExtensionId {
         &self.name
@@ -704,25 +665,6 @@ impl Extension {
         self.types.iter()
     }
 
-    /// Add a named static value to the extension.
-    pub fn add_value(
-        &mut self,
-        name: impl Into<ValueName>,
-        typed_value: ops::Value,
-    ) -> Result<&mut ExtensionValue, ExtensionBuildError> {
-        let extension_value = ExtensionValue {
-            extension: self.name.clone(),
-            name: name.into(),
-            typed_value,
-        };
-        match self.values.entry(extension_value.name.clone()) {
-            btree_map::Entry::Occupied(_) => {
-                Err(ExtensionBuildError::ValueExists(extension_value.name))
-            }
-            btree_map::Entry::Vacant(ve) => Ok(ve.insert(extension_value)),
-        }
-    }
-
     /// Instantiate an [`ExtensionOp`] which references an [`OpDef`] in this extension.
     pub fn instantiate_extension_op(
         &self,
@@ -783,9 +725,6 @@ pub enum ExtensionBuildError {
     /// Existing [`TypeDef`]
     #[error("Extension already has an type called {0}.")]
     TypeDefExists(TypeName),
-    /// Existing [`ExtensionValue`]
-    #[error("Extension already has an extension value called {0}.")]
-    ValueExists(ValueName),
 }
 
 /// A set of extensions identified by their unique [`ExtensionId`].

hugr-core/src/extension/const_fold.rs

@@ -2,18 +2,135 @@ use std::fmt::Formatter;
 
 use std::fmt::Debug;
 
+use crate::ops::constant::CustomConst;
+use crate::ops::constant::{OpaqueValue, Sum};
 use crate::ops::Value;
-use crate::types::TypeArg;
+use crate::types::{SumType, TypeArg};
+use crate::{IncomingPort, Node, OutgoingPort, PortIndex};
 
-use crate::IncomingPort;
-use crate::OutgoingPort;
+/// Representation of values used for constant folding.
+// Should we be non-exhaustive??
+// No point in parametrizing by HugrNode since then ConstFold would not be dyn/object-safe
+#[derive(Clone, Debug, PartialEq, Default)]
+pub enum FoldVal {
+    /// Value is unknown, must assume that it could be anything
+    #[default]
+    Unknown,
+    /// A variant of a [SumType]
+    Sum {
+        /// Which variant of the sum type this value is.
+        tag: usize,
+        /// Describes the type of the whole value.
+        // Can we deprecate this immediately? It is only for converting to Value
+        sum_type: SumType,
+        /// A value for each element (type) within the variant
+        items: Vec<FoldVal>,
+    },
+    /// A constant value defined by an extension
+    Extension(OpaqueValue),
+    /// A function pointer loaded from a [FuncDefn](crate::ops::FuncDefn) or `FuncDecl`
+    LoadedFunction(Node, Vec<TypeArg>), // Deliberately skipping Function(Box<Hugr>) ATM
+}
+
+impl<T> From<T> for FoldVal
+where
+    T: CustomConst,
+{
+    fn from(value: T) -> Self {
+        Self::Extension(value.into())
+    }
+}
 
-use crate::ops;
+impl FoldVal {
+    /// Returns a constant "false" value, i.e. the first variant of Sum((), ()).
+    pub const fn false_val() -> Self {
+        Self::Sum {
+            tag: 0,
+            sum_type: SumType::Unit { size: 2 },
+            items: vec![],
+        }
+    }
+
+    /// Returns a constant "true" value, i.e. the second variant of Sum((), ()).
+    pub const fn true_val() -> Self {
+        Self::Sum {
+            tag: 1,
+            sum_type: SumType::Unit { size: 2 },
+            items: vec![],
+        }
+    }
+
+    /// Returns a constant boolean - either [Self::false_val] or [Self::true_val]
+    pub const fn from_bool(b: bool) -> Self {
+        if b {
+            Self::true_val()
+        } else {
+            Self::false_val()
+        }
+    }
+
+    /// Extract the specified type of [CustomConst] fro this instance, if it is one
+    pub fn get_custom_value<T: CustomConst>(&self) -> Option<&T> {
+        let Self::Extension(e) = self else {
+            return None;
+        };
+        e.value().downcast_ref()
+    }
+}
+
+impl TryFrom<FoldVal> for Value {
+    type Error = Option<Node>;
+
+    fn try_from(value: FoldVal) -> Result<Self, Self::Error> {
+        match value {
+            FoldVal::Unknown => Err(None),
+            FoldVal::Sum {
+                tag,
+                sum_type,
+                items,
+            } => {
+                let values = items
+                    .into_iter()
+                    .map(Value::try_from)
+                    .collect::<Result<Vec<_>, _>>()?;
+                Ok(Value::Sum(Sum {
+                    tag,
+                    values,
+                    sum_type,
+                }))
+            }
+            FoldVal::Extension(e) => Ok(Value::Extension { e }),
+            FoldVal::LoadedFunction(node, _) => Err(Some(node)),
+        }
+    }
+}
+
+impl From<Value> for FoldVal {
+    fn from(value: Value) -> Self {
+        match value {
+            Value::Extension { e } => FoldVal::Extension(e),
+            #[allow(deprecated)] // remove when Value::Function removed
+            Value::Function { .. } => FoldVal::Unknown,
+            Value::Sum(Sum {
+                tag,
+                values,
+                sum_type,
+            }) => {
+                let items = values.into_iter().map(FoldVal::from).collect();
+                FoldVal::Sum {
+                    tag,
+                    sum_type,
+                    items,
+                }
+            }
+        }
+    }
+}
 
 /// Output of constant folding an operation, None indicates folding was either
 /// not possible or unsuccessful. An empty vector indicates folding was
 /// successful and no values are output.
-pub type ConstFoldResult = Option<Vec<(OutgoingPort, ops::Value)>>;
+pub type ConstFoldResult = Option<Vec<(OutgoingPort, Value)>>;
 
 /// Tag some output constants with [`OutgoingPort`] inferred from the ordering.
 pub fn fold_out_row(consts: impl IntoIterator<Item = Value>) -> ConstFoldResult {
@@ -27,9 +144,29 @@ pub fn fold_out_row(consts: impl IntoIterator<Item = Value>) -> ConstFoldResult
 
 /// Trait implemented by extension operations that can perform constant folding.
 pub trait ConstFold: Send + Sync {
+    /// Given type arguments `type_args` and [`FoldVal`]s for each input,
+    /// update the outputs (these will be initialized to [FoldVal::Unknown]).
+    ///
+    /// Defaults to calling [Self::fold] with those arguments that can be converted ---
+    /// [FoldVal::LoadedFunction]s will be lost as these are not representable as [Value]s.
+    fn fold2(&self, type_args: &[TypeArg], inputs: &[FoldVal], outputs: &mut [FoldVal]) {
+        let consts = inputs
+            .iter()
+            .cloned()
+            .enumerate()
+            .filter_map(|(p, fv)| Some((p.into(), fv.try_into().ok()?)))
+            .collect::<Vec<_>>();
+        #[allow(deprecated)] // remove this when fold is removed
+        let outs = self.fold(type_args, &consts);
+        for (p, v) in outs.unwrap_or_default() {
+            outputs[p.index()] = v.into();
+        }
+    }
+
     /// Given type arguments `type_args` and
     /// [`crate::ops::Const`] values for inputs at [`crate::IncomingPort`]s,
     /// try to evaluate the operation.
+    #[deprecated(note = "Use fold2")]
     fn fold(
         &self,
         type_args: &[TypeArg],

hugr-core/src/extension/op_def.rs

@@ -4,15 +4,16 @@ use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
 use std::sync::{Arc, Weak};
 
+use super::const_fold::FoldVal;
 use super::{
     ConstFold, ConstFoldResult, Extension, ExtensionBuildError, ExtensionId, ExtensionSet,
     SignatureError,
 };
 
-use crate::ops::{OpName, OpNameRef};
+use crate::ops::{OpName, OpNameRef, Value};
 use crate::types::type_param::{check_type_args, TypeArg, TypeParam};
 use crate::types::{FuncValueType, PolyFuncType, PolyFuncTypeRV, Signature};
-use crate::Hugr;
+use crate::{Hugr, IncomingPort};
 mod serialize_signature_func;
 
 /// Trait necessary for binary computations of OpDef signature
@@ -457,14 +458,30 @@ impl OpDef {
 
     /// Evaluate an instance of this [`OpDef`] defined by the `type_args`, given
     /// [`crate::ops::Const`] values for inputs at [`crate::IncomingPort`]s.
+    #[deprecated(note = "use constant_fold2")]
     pub fn constant_fold(
         &self,
         type_args: &[TypeArg],
-        consts: &[(crate::IncomingPort, crate::ops::Value)],
+        consts: &[(IncomingPort, Value)],
     ) -> ConstFoldResult {
+        #[allow(deprecated)] // we are in deprecated function, remove at same time
         (self.constant_folder.as_ref())?.fold(type_args, consts)
     }
 
+    /// Evaluate an instance of this [`OpDef`] defined by the `type_args`, given
+    /// [FoldVal] values for each input, and update the outputs, which should be
+    /// initialised to [FoldVal::Unknown].
+    pub fn constant_fold2(
+        &self,
+        type_args: &[TypeArg],
+        inputs: &[FoldVal],
+        outputs: &mut [FoldVal],
+    ) {
+        if let Some(cf) = self.constant_folder.as_ref() {
+            cf.fold2(type_args, inputs, outputs)
+        }
+    }
+
     /// Returns a reference to the signature function of this [`OpDef`].
     pub fn signature_func(&self) -> &SignatureFunc {
         &self.signature_func

hugr-core/src/extension/prelude/generic.rs

@@ -169,11 +169,14 @@ impl HasConcrete for LoadNatDef {
 mod tests {
     use crate::{
         builder::{inout_sig, DFGBuilder, Dataflow, DataflowHugr},
-        extension::prelude::{usize_t, ConstUsize},
-        ops::{constant, OpType},
+        extension::{
+            prelude::{usize_t, ConstUsize},
+            FoldVal,
+        },
+        ops::OpType,
         type_row,
         types::TypeArg,
-        HugrView, OutgoingPort,
+        HugrView,
     };
 
     use super::LoadNat;
@@ -209,10 +212,10 @@ mod tests {
 
         match optype {
             OpType::ExtensionOp(ext_op) => {
-                let result = ext_op.constant_fold(&[]);
-                let exp_port: OutgoingPort = 0.into();
-                let exp_val: constant::Value = ConstUsize::new(5).into();
-                assert_eq!(result, Some(vec![(exp_port, exp_val)]))
+                let mut out = [FoldVal::Unknown];
+                ext_op.constant_fold2(&[], &mut out);
+                let exp_val: FoldVal = ConstUsize::new(5).into();
+                assert_eq!(out, [exp_val])
             }
             _ => panic!(),
         }