diff --git a/hugr-core/src/hugr/rewrite/replace.rs b/hugr-core/src/hugr/rewrite/replace.rs index 80a847eea8..ca195e73d1 100644 --- a/hugr-core/src/hugr/rewrite/replace.rs +++ b/hugr-core/src/hugr/rewrite/replace.rs @@ -61,7 +61,7 @@ pub struct Replacement { /// and there would be no possible [Self::mu_inp], [Self::mu_out] or [Self::adoptions]. pub removal: Vec, /// A hugr (not necessarily valid, as it may be missing edges and/or nodes), whose root - /// is the same type as the root of [Self::replacement]. "G" in the spec. + /// is the same type as the common parent of all the nodes in [Self::removal]. "G" in the spec. pub replacement: Hugr, /// Describes how parts of the Hugr that would otherwise be removed should instead be preserved but /// with new parents amongst the newly-inserted nodes. This is a Map from container nodes in diff --git a/hugr-core/src/types.rs b/hugr-core/src/types.rs index 962e00876c..62ced85364 100644 --- a/hugr-core/src/types.rs +++ b/hugr-core/src/types.rs @@ -256,6 +256,15 @@ impl SumType { _ => None, } } + + /// Convert the SumType into its constituent rows (one per variant) + pub fn into_rows(self) -> impl Iterator { + match self { + SumType::Unit { size } => vec![TypeRowRV::new(); size as usize], + SumType::General { rows } => rows, + } + .into_iter() + } } impl From for TypeBase { diff --git a/hugr-passes/src/const_fold.rs b/hugr-passes/src/const_fold.rs index ce0f82476a..b54594d9d7 100644 --- a/hugr-passes/src/const_fold.rs +++ b/hugr-passes/src/const_fold.rs @@ -32,6 +32,7 @@ pub struct ConstantFoldPass { validation: ValidationLevel, allow_increase_termination: bool, inputs: HashMap, + entry_point: Option, } #[derive(Debug, Error)] @@ -71,6 +72,17 @@ impl ConstantFoldPass { self } + /// Sets the entry point for a [Module]-rooted Hugr, i.e. at which [FuncDefn] + /// child of the root the Hugr starts executing. If unspecified, for + /// [Module]-rooted Hugrs, the entire contents will be removed. + /// + /// [FuncDefn]: hugr_core::ops::OpType::FuncDefn + /// [Module]: hugr_core::ops::OpType::Module + pub fn with_module_entry_point(mut self, node: Node) -> Self { + self.entry_point = Some(node); + self + } + /// Run the Constant Folding pass. fn run_no_validate(&self, hugr: &mut impl HugrMut) -> Result<(), ConstFoldError> { let fresh_node = Node::from(portgraph::NodeIndex::new( @@ -88,8 +100,7 @@ impl ConstantFoldPass { }); let results = Machine::new(&hugr).run(ConstFoldContext(hugr), inputs); - let mut keep_nodes = HashSet::new(); - self.find_needed_nodes(&results, &mut keep_nodes); + let keep_nodes = self.find_needed_nodes(&results); let mb_root_inp = hugr.get_io(hugr.root()).map(|[i, _]| i); let remove_nodes = hugr @@ -145,17 +156,23 @@ impl ConstantFoldPass { fn find_needed_nodes( &self, results: &AnalysisResults, - needed: &mut HashSet, - ) { - let mut q = VecDeque::new(); + ) -> HashSet { let h = results.hugr(); - q.push_back(h.root()); + let mut needed = HashSet::new(); + let mut q = VecDeque::from_iter([h.root()]); + if let Some(entry) = self.entry_point { + assert!(h.get_parent(entry) == Some(h.root())); + assert!(h.get_optype(entry).is_func_defn()); + assert!(h.get_optype(h.root()).is_module()); + q.push_back(entry); + } + while let Some(n) = q.pop_front() { if !needed.insert(n) { continue; }; - if h.get_optype(n).is_cfg() { + if h.get_optype(n).is_cfg() || h.get_optype(n).is_conditional() { for bb in h.children(n) { //if results.bb_reachable(bb).unwrap() { // no, we'd need to patch up predicates q.push_back(bb); @@ -176,22 +193,19 @@ impl ConstantFoldPass { } // Also follow dataflow demand for (src, op) in h.all_linked_outputs(n) { - let needs_predecessor = match h.get_optype(src).port_kind(op).unwrap() { - EdgeKind::Value(_) => { - h.get_optype(src).is_load_constant() - || results - .try_read_wire_concrete::(Wire::new(src, op)) - .is_err() - } - EdgeKind::StateOrder | EdgeKind::Const(_) | EdgeKind::Function(_) => true, - EdgeKind::ControlFlow => false, // we always include all children of a CFG above - _ => true, // needed as EdgeKind non-exhaustive; not knowing what it is, assume the worst - }; - if needs_predecessor { - q.push_back(src); + if matches!(h.get_optype(src).port_kind(op).unwrap(), EdgeKind::Value(_)) + && results + .try_read_wire_concrete::(Wire::new(src, op)) + .is_ok() + && !h.get_optype(src).is_load_constant() + { + continue; } + // All other edge types --> we need the predecessors (even if included already e.g. ControlFlow) + q.push_back(src); } } + needed } } diff --git a/hugr-passes/src/lib.rs b/hugr-passes/src/lib.rs index ffc739933f..2e24ee1f87 100644 --- a/hugr-passes/src/lib.rs +++ b/hugr-passes/src/lib.rs @@ -27,6 +27,8 @@ pub use monomorphize::monomorphize; pub use monomorphize::{MonomorphizeError, MonomorphizePass}; pub mod nest_cfgs; pub mod non_local; +mod static_eval; +pub use static_eval::static_eval; pub mod validation; pub use force_order::{force_order, force_order_by_key}; pub use lower::{lower_ops, replace_many_ops}; diff --git a/hugr-passes/src/static_eval.rs b/hugr-passes/src/static_eval.rs new file mode 100644 index 0000000000..7537fef043 --- /dev/null +++ b/hugr-passes/src/static_eval.rs @@ -0,0 +1,391 @@ +//! Static evaluation of Hugrs, i.e. down to a [Value] (given inputs). +//! Some of this logic might be generalizable to cases where we cannot deduce a +//! unique [Value], and thus integrated into [constant folding](super::const_fold), +//! but the API is useful, and it seems likely that some of the transforms +//! will not be worth performing if a single-[Value] is not wanted and/or achievable. + +use std::collections::HashMap; + +use hugr_core::hugr::hugrmut::HugrMut; +use hugr_core::hugr::internal::HugrMutInternals; +use hugr_core::hugr::rewrite::inline_dfg::InlineDFG; +use hugr_core::hugr::rewrite::replace::{NewEdgeKind, NewEdgeSpec, Replacement}; +use hugr_core::hugr::views::{DescendantsGraph, ExtractHugr, HierarchyView}; +use hugr_core::ops::constant::Sum; +use hugr_core::ops::handle::{FuncID, TailLoopID}; +use hugr_core::ops::{ + Case, Conditional, Const, DataflowOpTrait, DataflowParent, Input, LoadConstant, OpType, Output, + Value, DFG, +}; +use hugr_core::types::{Signature, TypeEnum}; +use hugr_core::{Direction, Hugr, HugrView, Node, PortIndex}; + +use crate::const_fold::ConstantFoldPass; + +pub fn static_eval(mut h: Hugr, entry: Option) -> Option> { + // TODO: allow inputs to be specified + let mut cp = ConstantFoldPass::default(); + if let Some(ep) = entry { + cp = cp.with_module_entry_point(ep); + }; + let start = entry.unwrap_or(h.root()); + 'reanalyse: loop { + cp.run(&mut h).unwrap(); + eprintln!("****Constant folding produced {}", h.mermaid_string()); + loop { + let mut need_reanalyse = false; + let mut need_scan = false; + for n in h.children(start).collect::>() { + match h.get_optype(n) { + OpType::Conditional(_) => { + need_scan |= conditional_to_dfg(&mut h, n).is_some(); + // will inline on next iter. PERF: inline it now + } + OpType::DFG(_) => { + h.apply_rewrite(InlineDFG(n.into())).unwrap(); + need_scan = true; + } + OpType::TailLoop(_) => { + // Even if no inputs are constants (e.g. they are partial sums with multiple tags), + // peeling increases precision so could be beneficial. Note that it moves the TailLoop + // inside a Conditional, and we only peel TailLoops at the top level of the Hugr. + peel_tailloop(&mut h, n); + need_reanalyse = true; + } + OpType::CallIndirect(_) => { + let (called, _) = h.single_linked_output(n, 0).unwrap(); + match h.get_optype(called) { + OpType::LoadConstant(_) => { + // TODO: Inline called Hugr into DFG + need_reanalyse = true; + } + OpType::LoadFunction(_) => { + // TODO: Convert to Call + need_reanalyse = true + } + _ => (), + } + } + OpType::Call(_) => { + // Even if no inputs are constants (e.g. they are partial sums with multiple tags), + // inlining *could* (maybe) be beneficial. + // Note we are only doing this at the top level of the Hugr! + need_reanalyse |= inline_call(&mut h, n).is_some(); + } + OpType::CFG(_) => { + // TODO: if entry node has in-edges (i.e. from other blocks) -> peel, set precision_improved=True + // else if entry node is exit block, elide CFG; + // else if entry-node predicate is constant -> move contents of entry node outside CFG, make selected successor be the new entry block, set need_scan=True + } + _ => (), + } + } + eprintln!( + "Scan of top level ({} reanalysis, {} rescan) produced {}", + if need_reanalyse { "needs" } else { "no" }, + if need_scan { "needs" } else { "no" }, + h.mermaid_string() + ); + h.validate().unwrap(); + if need_reanalyse { + break; + }; + if !need_scan { + break 'reanalyse; // done + }; + } + } + + let [_, out] = h.get_io(start).unwrap(); + h.signature(out) + .unwrap() + .input_ports() + .map(|p| { + let (src_node, _) = h.single_linked_output(out, p)?; + h.get_optype(src_node).as_load_constant()?; + let cst = h.get_optype(h.static_source(src_node)?).as_const()?; + Some(cst.value().clone()) + }) + .collect() +} + +fn conditional_to_dfg(h: &mut impl HugrMut, cond: Node) -> Option<()> { + let (pred, _) = h.single_linked_output(cond, 0).unwrap(); + h.get_optype(pred).as_load_constant()?; + let cst_node = h.static_source(pred).unwrap(); + let Value::Sum(Sum { tag, values, .. }) = h.get_optype(cst_node).as_const().unwrap().value() + else { + panic!("Conditional input was not a Sum") + }; + let case_node = h.children(cond).nth(*tag).unwrap(); + let signature = h.get_optype(case_node).as_case().unwrap().signature.clone(); + + let mut replacement = Hugr::new(h.get_optype(h.get_parent(cond).unwrap()).clone()); + let dfg = replacement.add_node_with_parent(replacement.root(), DFG { signature }); + for (i, v) in values.iter().enumerate() { + let cst = replacement.add_node_with_parent(replacement.root(), Const::new(v.clone())); + let datatype = v.get_type(); + let lcst = replacement.add_node_after(cst, LoadConstant { datatype }); + replacement.connect(cst, 0, lcst, 0); + replacement.connect(lcst, 0, dfg, i); + } + let mut removal = vec![cond]; + if h.static_targets(pred).map_or(0, Iterator::count) == 1 { + // Also remove the original (Sum) constant - we could leave this for later DCE? + removal.push(pred); + if h.static_targets(cst_node).map_or(0, Iterator::count) == 1 { + removal.push(cst_node); + } + } + h.apply_rewrite(Replacement { + removal, + replacement, + adoptions: HashMap::from([(dfg, case_node)]), + mu_inp: make_mu(h, cond, dfg, Direction::Incoming, true) // predicate dealt with above + .map(|mut e| { + if let NewEdgeKind::Value { tgt_pos, .. } = &mut e.kind { + *tgt_pos = (tgt_pos.index() + values.len() - 1).into(); + } + e + }) + .collect(), + mu_new: vec![], + mu_out: make_mu(h, cond, dfg, Direction::Outgoing, false).collect(), + }) + .unwrap(); + Some(()) +} + +fn make_mu( + h: &impl HugrView, + old: Node, + new: Node, + dir: Direction, + skip_first: bool, +) -> impl Iterator + '_ { + let ord_port = h.get_optype(old).other_port(dir); + h.node_ports(old, dir) + .skip(skip_first as usize) + .flat_map(move |pos| { + h.linked_ports(old, pos).map(move |(node, nodp)| { + let (src, tgt) = if dir == Direction::Outgoing { + (new, node) + } else { + (node, new) + }; + let kind = if Some(pos) == ord_port { + NewEdgeKind::Order + } else { + NewEdgeKind::Value { + src_pos: pos.as_outgoing().or_else(|_| nodp.as_outgoing()).unwrap(), + tgt_pos: pos.as_incoming().or_else(|_| nodp.as_incoming()).unwrap(), + } + }; + NewEdgeSpec { src, tgt, kind } + }) + }) +} + +fn inline_call(h: &mut impl HugrMut, call: Node) -> Option<()> { + let orig_func = h.static_source(call).unwrap(); + let function = DescendantsGraph::>::try_new(&h, orig_func).ok()?; + // Ideally we'd like the following to preserve uses from within "function" of Consts outside + // the function, but (see https://github.com/CQCL/hugr/discussions/1642) this probably won't happen at the moment - TODO XXX FIXME + let mut func = function.extract_hugr(); + let recursive_calls = func + .static_targets(func.root()) + .unwrap() + .collect::>(); + let new_op = OpType::from(DFG { + signature: func + .root_type() + .as_func_defn() + .unwrap() + .inner_signature() + .into_owned(), + }); + let (in_ports, out_ports) = (new_op.input_count(), new_op.output_count()); + func.replace_op(func.root(), new_op).unwrap(); + func.set_num_ports(func.root(), in_ports as _, out_ports as _); + let func_copy = h.insert_hugr(h.get_parent(call).unwrap(), func); + for (rc, p) in recursive_calls.into_iter() { + let call_node = func_copy.node_map.get(&rc).unwrap(); + h.disconnect(*call_node, p); + h.connect(orig_func, 0, *call_node, p); + } + let func_copy = func_copy.new_root; + let new_connections = h + .all_linked_outputs(call) + .filter(|(n, _)| *n != orig_func) + .enumerate() + .map(|(tgt_port, (src, src_port))| (src, src_port, func_copy, tgt_port.into())) + .chain(h.node_outputs(call).flat_map(|src_port| { + h.linked_inputs(call, src_port) + .map(move |(tgt, tgt_port)| (func_copy, src_port, tgt, tgt_port)) + })) + .collect::>(); + h.remove_node(call); + for (src_node, src_port, tgt_node, tgt_port) in new_connections { + h.connect(src_node, src_port, tgt_node, tgt_port); + } + Some(()) +} + +fn peel_tailloop(h: &mut impl HugrMut, tl: Node) { + // TODO: copy body of loop into DFG (dup loop, change container type, output Sum) + // TODO: change constant into elements + // TODO: nest existing loop inside conditional testing output of DFG. + let tl_desc = h.get_optype(tl).as_tail_loop().unwrap(); + let mut replacement = Hugr::new(h.get_optype(h.get_parent(tl).unwrap()).clone()); + let first_iter = replacement + .insert_from_view( + replacement.root(), + &DescendantsGraph::::try_new(h, tl).unwrap(), + ) + .new_root; + + let signature = tl_desc.inner_signature().into_owned(); + let outer_sig = tl_desc.signature().into_owned(); + let cond = { + // Converts the result of the first iteration, i.e. results *inside* the TailLoop, + // to the result of the whole TailLoop + let mut iter_result = signature.output.iter(); + let TypeEnum::Sum(st) = iter_result.next().unwrap().as_type_enum() else { + panic!("First output of loop body was not predicate") + }; + Conditional { + // The loop's control predicate cannot actually contain any Row Variables + sum_rows: st + .clone() + .into_rows() + .map(|trv| trv.try_into().unwrap()) + .collect(), + other_inputs: iter_result.cloned().collect::>().into(), + outputs: outer_sig.output.clone(), + extension_delta: signature.runtime_reqs.clone(), + } + }; + debug_assert_eq!(cond.signature().input, signature.output); + let cond = replacement.add_node_after(first_iter, cond); + let dfg = OpType::from(DFG { signature }); + let (in_count, out_count) = (dfg.input_count(), dfg.output_count()); + replacement.replace_op(first_iter, dfg).unwrap(); + replacement.set_num_ports(first_iter, in_count, out_count); + fn wire_all(h: &mut Hugr, from: Node, to: Node) { + for p in h.node_outputs(from).collect::>() { + h.connect(from, p, to, p.index()); + } + } + wire_all(&mut replacement, first_iter, cond); + // Continue variant: the original TailLoop will go in here + let cont = replacement.add_node_with_parent( + cond, + Case { + signature: outer_sig.clone(), + }, + ); + let inp = replacement.add_node_with_parent( + cont, + Input { + types: outer_sig.input.clone(), + }, + ); + let oup = replacement.add_node_with_parent( + cont, + Output { + types: outer_sig.output.clone(), + }, + ); + let new_tl = replacement.add_node_with_parent(cont, tl_desc.clone()); + wire_all(&mut replacement, inp, new_tl); + wire_all(&mut replacement, new_tl, oup); + + // Break variant + let brk = replacement.add_node_after( + cont, + Case { + signature: Signature::new_endo(outer_sig.output.clone()), + }, + ); + let inp = replacement.add_node_with_parent( + brk, + Input { + types: outer_sig.output.clone(), + }, + ); + let oup = replacement.add_node_with_parent( + brk, + Output { + types: outer_sig.output, + }, + ); + wire_all(&mut replacement, inp, oup); + + h.apply_rewrite(Replacement { + removal: vec![tl], + replacement, + adoptions: HashMap::from([(new_tl, tl)]), + mu_inp: make_mu(h, tl, first_iter, Direction::Incoming, false).collect(), + mu_out: make_mu(h, tl, cond, Direction::Outgoing, false).collect(), + mu_new: vec![], + }) + .unwrap(); +} + +#[cfg(test)] +mod test { + use std::{fs::File, io::BufReader}; + + use hugr_core::std_extensions::{arithmetic::int_types::ConstInt, STD_REG}; + use hugr_core::HugrView; + use hugr_core::{ops::Value, Hugr}; + use itertools::Itertools; + + use super::static_eval; + + #[test] + fn recursive_fibonacci() { + let h = Hugr::load_json( + BufReader::new(File::open("/Users/alanlawrence/fibonacci_hugr.json").unwrap()), + &STD_REG, + ) + .unwrap(); + let main = h + .children(h.root()) + .filter(|n| { + h.get_optype(*n) + .as_func_defn() + .is_some_and(|f| f.name == "main") + }) + .exactly_one() + .ok() + .unwrap(); + assert_eq!( + static_eval(h, Some(main)), + Some(vec![Value::extension(ConstInt::new_u(5, 8).unwrap())]) + ); + } + + #[test] + fn iterative_tailloop() { + let h = Hugr::load_json( + BufReader::new(File::open("/Users/alanlawrence/factorial_hugr.json").unwrap()), + &STD_REG, + ) + .unwrap(); + let main = h + .children(h.root()) + .filter(|n| { + h.get_optype(*n) + .as_func_defn() + .is_some_and(|f| f.name == "main") + }) + .exactly_one() + .ok() + .unwrap(); + assert_eq!( + static_eval(h, Some(main)), + Some(vec![Value::extension(ConstInt::new_u(5, 120).unwrap())]) + ); + } +} diff --git a/hugr-py/src/hugr/std/int.py b/hugr-py/src/hugr/std/int.py index 8652cc1359..92a2380aee 100644 --- a/hugr-py/src/hugr/std/int.py +++ b/hugr-py/src/hugr/std/int.py @@ -106,3 +106,97 @@ def __call__(self, a: ComWire, b: ComWire) -> Command: #: DivMod operation. DivMod = _DivModDef() + + +@dataclass(frozen=True) +class _ILtUDef(RegisteredOp): + """Integer less than (unsigned).""" + + width: int = 5 + const_op_def: ClassVar[ext.OpDef] = INT_OPS_EXTENSION.operations["ilt_u"] + + def type_args(self) -> list[tys.TypeArg]: + return [tys.BoundedNatArg(n=self.width)] + + def cached_signature(self) -> tys.FunctionType | None: + row: list[tys.Type] = [int_t(self.width)] * 2 + return tys.FunctionType(row, [tys.Bool], runtime_reqs=[INT_OPS_EXTENSION.name]) + + def __call__(self, a: ComWire) -> Command: + return DataflowOp.__call__(self, a) + + +#: IntLessThan (unsigned) operation +ILtU = _ILtUDef() + + +@dataclass(frozen=True) +class _IMulDef(RegisteredOp): + """Integer multiply.""" + + width: int = 5 + const_op_def: ClassVar[ext.OpDef] = INT_OPS_EXTENSION.operations["imul"] + + def type_args(self) -> list[tys.TypeArg]: + return [tys.BoundedNatArg(n=self.width)] + + def cached_signature(self) -> tys.FunctionType | None: + row: list[tys.Type] = [int_t(self.width)] * 2 + return tys.FunctionType( + row, [int_t(self.width)], runtime_reqs=[INT_OPS_EXTENSION.name] + ) + + def __call__(self, a: ComWire) -> Command: + return DataflowOp.__call__(self, a) + + +#: IMul operation +IMul = _IMulDef() + + +@dataclass(frozen=True) +class _ISubDef(RegisteredOp): + """Integer subtract.""" + + width: int = 5 + const_op_def: ClassVar[ext.OpDef] = INT_OPS_EXTENSION.operations["isub"] + + def type_args(self) -> list[tys.TypeArg]: + return [tys.BoundedNatArg(n=self.width)] + + def cached_signature(self) -> tys.FunctionType | None: + row: list[tys.Type] = [int_t(self.width)] * 2 + return tys.FunctionType( + row, [int_t(self.width)], runtime_reqs=[INT_OPS_EXTENSION.name] + ) + + def __call__(self, a: ComWire) -> Command: + return DataflowOp.__call__(self, a) + + +#: ISub operation +ISub = _ISubDef() + + +@dataclass(frozen=True) +class _IAddDef(RegisteredOp): + """Integer add.""" + + width: int = 5 + const_op_def: ClassVar[ext.OpDef] = INT_OPS_EXTENSION.operations["iadd"] + + def type_args(self) -> list[tys.TypeArg]: + return [tys.BoundedNatArg(n=self.width)] + + def cached_signature(self) -> tys.FunctionType | None: + row: list[tys.Type] = [int_t(self.width)] * 2 + return tys.FunctionType( + row, [int_t(self.width)], runtime_reqs=[INT_OPS_EXTENSION.name] + ) + + def __call__(self, a: ComWire) -> Command: + return DataflowOp.__call__(self, a) + + +#: IAdd operation +IAdd = _IAddDef() diff --git a/hugr-py/tests/test_cond_loop.py b/hugr-py/tests/test_cond_loop.py index ae84f8cb83..9a4d38e6ac 100644 --- a/hugr-py/tests/test_cond_loop.py +++ b/hugr-py/tests/test_cond_loop.py @@ -3,8 +3,9 @@ from hugr import ops, tys, val from hugr.build.cond_loop import Conditional, ConditionalError, TailLoop from hugr.build.dfg import Dfg +from hugr.build.function import Module from hugr.package import Package -from hugr.std.int import INT_T, IntVal +from hugr.std.int import INT_T, ILtU, IMul, IntVal, ISub from .conftest import QUANTUM_EXT, H, Measure, validate @@ -144,3 +145,28 @@ def test_conditional_bug() -> None: with cond.add_case(0) as case: case.set_outputs() validate(cond.hugr) + + +def test_iterative_factorial() -> None: + mod = Module() + + fac = mod.define_function("factorial", [INT_T], [INT_T]) + loop_t = tys.Either([INT_T], []) + with fac.add_tail_loop(fac.inputs(), [fac.load(IntVal(1))]) as tl: + one = tl.load(IntVal(1)) + with tl.add_if(tl.add_op(ILtU, tl.input_node[0], one)) as if_: + if_.set_outputs(if_.add(ops.Break(loop_t)()), tl.input_node[1]) + with if_.add_else() as else_: + i2 = else_.add_op(ISub, tl.input_node[0], one) + f2 = else_.add_op(IMul, tl.input_node[0], tl.input_node[1]) + else_.set_outputs(else_.add(ops.Continue(loop_t)(i2)), f2) + tl.set_loop_outputs(*else_.conditional_node.outputs()) + fac.set_outputs(tl) + + main = mod.define_function("main", [], [INT_T]) + main.set_outputs(main.call(fac, main.load(IntVal(5)))) + + validate(mod.hugr) + + with open("/Users/alanlawrence/factorial_hugr.json", "w") as f: + f.write(mod.hugr.to_json()) diff --git a/hugr-py/tests/test_hugr_build.py b/hugr-py/tests/test_hugr_build.py index 46e0d37d09..da2fb155cf 100644 --- a/hugr-py/tests/test_hugr_build.py +++ b/hugr-py/tests/test_hugr_build.py @@ -10,7 +10,7 @@ from hugr.hugr import Hugr from hugr.hugr.node_port import Node, _SubPort from hugr.ops import NoConcreteFunc -from hugr.std.int import INT_T, DivMod, IntVal +from hugr.std.int import INT_T, DivMod, IAdd, ILtU, IntVal, ISub from hugr.std.logic import Not from .conftest import validate @@ -297,6 +297,31 @@ def test_invalid_recursive_function() -> None: f_recursive.set_outputs(f_recursive.input_node[0]) +def test_recursive_fibonacci() -> None: + mod = Module() + + fib = mod.define_function("fibonacci", [INT_T], [INT_T]) + one = fib.load(IntVal(1)) + two = fib.load(IntVal(2)) + pred = fib.add_op(ILtU, fib.input_node[0], two) + cond = fib.add_conditional(pred) + with cond.add_case(0) as f: + r1 = f.call(fib, f.add_op(ISub, fib.input_node[0], one)) + r2 = f.call(fib, f.add_op(ISub, fib.input_node[0], two)) + f.set_outputs(f.add_op(IAdd, r1, r2)) + with cond.add_case(1) as t: + t.set_outputs(t.load(IntVal(1))) + fib.set_outputs(*cond.outputs()) + + main = mod.define_function("main", [], [INT_T]) + main.set_outputs(main.call(fib, main.load(IntVal(5)))) + + validate(mod.hugr) + + with open("/Users/alanlawrence/fibonacci_hugr.json", "w") as f: + f.write(mod.hugr.to_json()) + + def test_higher_order() -> None: noop_fn = Dfg(tys.Qubit) noop_fn.set_outputs(noop_fn.add(ops.Noop()(noop_fn.input_node[0])))