handwired_rs/
inline.rs

1//! Inlining one circuit into another, ported from handwired's `inlineCircuit`.
2
3use std::collections::HashMap;
4
5use crate::circuit::{Circuit, Comparison, Macro, OpKind};
6use crate::wire::{Lin, Term, Wire};
7
8impl Circuit {
9    /// Rebuild every operation of `source` in this circuit, with source cell `k` read from `source_wires[k]`, and
10    /// return where each of `source`'s wires landed, indexed by its wire. The operations go through this circuit's
11    /// builder, so an expression that already exists here is shared, e.g. a sum both circuits compute. Declared
12    /// ranges, grids and exact bits of `source`'s cells land on the mapped wires, and rounding points, comparison
13    /// bits and one-hot groups carry over. With `allow_existing_readers`, those declarations may narrow wires that
14    /// operations here already read (see [`assume_range_allowing_readers`](Self::assume_range_allowing_readers)).
15    pub fn inline_circuit(
16        &mut self,
17        source: &Circuit,
18        source_wires: &[Wire],
19        allow_existing_readers: bool,
20    ) -> Vec<Wire> {
21        assert_eq!(source_wires.len(), source.source_size(), "inlined source mapping has the wrong length");
22        let mut map: Vec<Wire> = source_wires.to_vec();
23        let declare_range = |circuit: &mut Circuit, wire: Wire, min: f64, max: f64| {
24            if allow_existing_readers {
25                circuit.assume_range_allowing_readers(wire, min, max);
26            } else {
27                circuit.assume_range(wire, min, max);
28            }
29        };
30        let declare_grid = |circuit: &mut Circuit, wire: Wire, step: f64| {
31            if allow_existing_readers {
32                circuit.assume_grid_allowing_readers(wire, step);
33            } else {
34                circuit.assume_grid(wire, step);
35            }
36        };
37        for (cell, &to) in source_wires.iter().enumerate() {
38            let from = Wire::from_index(cell);
39            let range = source.range(from);
40            declare_range(self, to, range.min, range.max);
41            if let Some(grid) = source.grid(from) {
42                declare_grid(self, to, grid);
43            }
44            if source.is_exact_bit(from) {
45                declare_range(self, to, 0.0, 1.0);
46                self.mark_exact_bit(to);
47            }
48        }
49        // Each distinct table is mapped once, so reads of one large table share one mapped array.
50        let mut mapped_tables: HashMap<usize, Vec<Wire>> = HashMap::new();
51        for (operation_index, operation) in source.operations().iter().enumerate() {
52            let map_terms = |terms: &[Term]| -> Vec<Term> {
53                terms.iter().map(|&(wire, coefficient)| (map[wire.index()], coefficient)).collect()
54            };
55            let wire = match (operation.macro_op, operation.kind) {
56                (Some(Macro::Multiply { left, right }), _) => self.multiply(map[left.index()], map[right.index()]),
57                (Some(Macro::Select { gate, off, on }), _) => {
58                    self.select(map[gate.index()], map[off.index()], map[on.index()])
59                }
60                (None, OpKind::Reglu) => self.append(
61                    OpKind::Reglu,
62                    Lin::new(map_terms(&operation.left_terms), operation.left_bias),
63                    Lin::new(map_terms(&operation.gate_terms), operation.gate_bias),
64                    operation.preserve_rounding,
65                    None,
66                ),
67                (None, OpKind::Index) => {
68                    let read = operation.index_read.as_ref().unwrap();
69                    let candidates = mapped_tables
70                        .entry(read.table)
71                        .or_insert_with(|| {
72                            source.tables[read.table].candidates.iter().map(|c| map[c.index()]).collect()
73                        })
74                        .clone();
75                    // The accepted index range carries over, e.g. one declared on a small affine index.
76                    self.index_read_affine(
77                        Lin::new(map_terms(&operation.gate_terms), operation.gate_bias),
78                        &candidates,
79                        read.below,
80                        read.above,
81                        Some(read.index_range),
82                    )
83                }
84                // emit_compare's atom on grid operands reads its score terms directly.
85                (None, OpKind::Boolean) if operation.preserve_rounding => self.append(
86                    OpKind::Boolean,
87                    Lin::new(map_terms(&operation.left_terms), operation.left_bias),
88                    Lin::default(),
89                    true,
90                    None,
91                ),
92                (None, kind) => {
93                    let form = Lin::new(map_terms(&operation.left_terms), operation.left_bias);
94                    let value = if operation.preserve_rounding { self.round_affine(form) } else { self.affine(form) };
95                    if kind == OpKind::Boolean { self.boolean(value) } else { value }
96                }
97            };
98            let from = Wire::from_index(source.source_size() + operation_index);
99            let range = source.range(from);
100            declare_range(self, wire, range.min, range.max);
101            if let Some(grid) = source.grid(from) {
102                declare_grid(self, wire, grid);
103            }
104            // A marked operation can fold into a source cell here, which has no operation to record it on.
105            if let Some(index) = wire.index().checked_sub(self.source_size()) {
106                if let Some(kind) = operation.rounding_point {
107                    self.record_rounding_point(wire, kind);
108                }
109                // The mapped operation computes the same values from the mapped operands, so it still holds 0 or 1.
110                if operation.comparison_bit {
111                    self.operations[index].comparison_bit = true;
112                }
113                if let Some(comparison) = &operation.comparison {
114                    self.operations[index].comparison = Some(Comparison {
115                        terms: comparison
116                            .terms
117                            .iter()
118                            .map(|&(wire, coefficient)| (map[wire.index()], coefficient))
119                            .collect(),
120                        bias: comparison.bias,
121                        equal: comparison.equal,
122                    });
123                }
124            }
125            map.push(wire);
126        }
127        // One-hot groups are about runtime values, which the mapped wires still carry.
128        for group in source.one_hot_groups() {
129            let members: Vec<Wire> = group.members.iter().map(|member| map[member.index()]).collect();
130            if allow_existing_readers {
131                self.assume_one_hot_allowing_readers(&members, group.allow_all_zero);
132            } else {
133                self.assume_one_hot(&members, group.allow_all_zero);
134            }
135        }
136        map
137    }
138}
handwired_rs/inline.rs

handwired_rs/
inline.rs
