Include folding of extension ops!

CQCL · Aug 28, 2024 · 1594e6f · 1594e6f
1 parent ada1884
commit 1594e6f
Showing 1 changed file with 15 additions and 3 deletions.
diff --git a/hugr-passes/src/const_fold2/datalog/context.rs b/hugr-passes/src/const_fold2/datalog/context.rs
@@ -2,10 +2,10 @@ use std::hash::{Hash, Hasher};
 use std::ops::Deref;
 use std::sync::Arc;
 
-use hugr_core::ops::OpType;
-use hugr_core::{Hugr, HugrView, Node};
+use hugr_core::ops::{CustomOp, DataflowOpTrait, OpType};
+use hugr_core::{Hugr, HugrView, IncomingPort, Node, PortIndex};
 
-use crate::const_fold2::value_handle::ValueHandle;
+use crate::const_fold2::value_handle::{ValueHandle, ValueKey};
 
 use super::utils::value_outputs;
 use super::{DFContext, PartialValue};
@@ -82,6 +82,18 @@ impl<H: HugrView> DFContext<ValueHandle> for DataflowContext<H> {
                 )
                 .into()])
             }
+            OpType::CustomOp(CustomOp::Extension(op)) => {
+                let sig = op.signature();
+                let known_ins = sig.input_types().into_iter().enumerate().zip(ins.iter()).filter_map(|((i,ty), pv)|
+                    pv.clone().try_into_value(ty).map(|v| (IncomingPort::from(i), v)).ok())
+                    .collect::<Vec<_>>();
+                let outs = op.constant_fold(&known_ins)?;
+                let mut res = vec![PartialValue::bottom(); sig.output_count()];
+                for (op, v) in outs {
+                    res[op.index()] = ValueHandle::new(ValueKey::Node(n), Arc::new(v)).into()
+                }
+                Some(res)
+            }
             // This is a bit awkward - these seem very generic across all DFContexts
             OpType::MakeTuple(_) => Some(vec![PartialValue::variant(0, ins.into_iter().cloned())]),
             OpType::UnpackTuple(_) => {