risc0_zkp/
layout.rs

1// Copyright 2024 RISC Zero, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Utilities for accessing buffers of circuits symbolically.
16
17use alloc::{collections::BTreeMap, string::String, vec::Vec};
18#[cfg(not(feature = "std"))]
19use alloc::{format, string::ToString};
20use anyhow::Result;
21use core::fmt::{Debug, Formatter, Write};
22pub use paste::paste;
23use risc0_core::field::Elem;
24
25/// A reference to a register in a layout
26pub struct Reg {
27    pub offset: usize,
28}
29
30impl Component for Reg {
31    fn walk<V: Visitor>(&self, v: &mut V) -> core::fmt::Result {
32        v.visit_reg(self.offset)
33    }
34    fn ty_name(&self) -> &'static str {
35        "reg"
36    }
37}
38
39/// A component within a circuit execution trace buffer.  Users should
40/// not use this directly; it is only "pub" so it can be used by the
41/// generated layout code.
42#[doc(hidden)]
43pub trait Component {
44    fn walk<V: Visitor>(&self, v: &mut V) -> core::fmt::Result;
45    fn ty_name(&self) -> &'static str;
46}
47
48impl<C: Component, const N: usize> Component for [&C; N] {
49    fn walk<V: Visitor>(&self, v: &mut V) -> core::fmt::Result {
50        for (n, elem) in self.iter().enumerate() {
51            let name = alloc::format!("[{n}]");
52            v.visit_component(&name, *elem)?;
53        }
54        Ok(())
55    }
56    fn ty_name(&self) -> &'static str {
57        "array"
58    }
59}
60
61/// A visitor that visits components in a component tree. Users should
62/// not use this directly; it is only "pub" so it can be used by the
63/// generated layout code.
64#[doc(hidden)]
65pub trait Visitor {
66    fn visit_component(&mut self, name: &str, component: &impl Component) -> core::fmt::Result;
67    fn visit_reg(&mut self, offset: usize) -> core::fmt::Result;
68}
69
70/// Represents a section of a component tree that's present in a buffer.
71#[derive(Clone, Copy)]
72pub struct Tree<'a, E: Elem + Into<u32>, C: Component> {
73    pub buf: &'a [E],
74    pub component: &'a C,
75}
76
77impl<'a, E: Elem + Into<u32>, C: Component> Tree<'a, E, C> {
78    /// Returns the component tree rooted at the given component.
79    pub fn new(buf: &'a [E], component: &'a C) -> Self {
80        Tree { buf, component }
81    }
82
83    /// Descends into a subtree.  The given function should return the
84    /// requested layout based when passed this tree's layout.
85    pub fn map<SUBC: Component>(&self, f: impl FnOnce(&'a C) -> &'a SUBC) -> Tree<'a, E, SUBC> {
86        Tree {
87            buf: self.buf,
88            component: f(self.component),
89        }
90    }
91
92    /// Interprets the contents of this tree as a list of u32s.
93    pub fn get_u32s(&self) -> Result<Vec<u32>> {
94        let mut gather = TreeGather::new(self.buf);
95        self.component
96            .walk(&mut gather)
97            .map_err(anyhow::Error::msg)?;
98        Ok(gather.vals)
99    }
100
101    /// Interprets the contents of this tree as a list of bytes.
102    pub fn get_bytes(&self) -> Result<Vec<u8>> {
103        self.get_u32s()?
104            .into_iter()
105            .map(|val| u8::try_from(val).map_err(anyhow::Error::msg))
106            .collect()
107    }
108
109    /// Returns the contents of this tree as a u32; elements are expected to be
110    /// 4 bytes.
111    pub fn get_u32_from_bytes(&self) -> Result<u32> {
112        Ok(u32::from_le_bytes(
113            self.get_bytes()?
114                .as_slice()
115                .try_into()
116                .map_err(anyhow::Error::msg)?,
117        ))
118    }
119
120    /// Returns the contents of this tree, which should have a single element, as a u32.
121    pub fn get_u32_from_elem(&self) -> Result<u32> {
122        let u32s = self.get_u32s()?;
123        assert_eq!(u32s.len(), 1, "Expecting only a single u32 in {:?}", self);
124        Ok(u32s[0])
125    }
126}
127
128/// Retrieves the given register from a buffer and returns it as a field element
129pub fn get_elem<'a, E>(buf: &'a [E], reg: &Reg) -> &'a E {
130    buf.get(reg.offset).expect("Invalid offset in layout")
131}
132
133/// Retrieves the given register from a buffer and interprets it as a u32
134pub fn get_u32<E: Copy + Into<u32>>(buf: &[E], reg: &Reg) -> u32 {
135    (*get_elem(buf, reg)).into()
136}
137
138impl<'a, E: Elem + Into<u32>, C: Component> Debug for Tree<'a, E, C> {
139    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
140        let mut p = TreePrinter::new(self.buf, Vec::new(), "top");
141        self.component.walk(&mut p)?;
142        for line in p.lines {
143            writeln!(f, "{line}")?;
144        }
145        Ok(())
146    }
147}
148
149// Collects a tree dump as a set of lines.  Lists identical sections as a
150// reference to the first one seen.
151struct TreePrinter<'a, E: Elem> {
152    buf: &'a [E],
153    lines: Vec<String>,
154    item_count: usize,
155    seen: BTreeMap<Vec<String>, String>,
156    path: Vec<&'a str>,
157}
158
159impl<'a, E: Elem> TreePrinter<'a, E> {
160    fn new(buf: &'a [E], mut path: Vec<&'a str>, name: &'a str) -> Self {
161        path.push(name);
162        TreePrinter {
163            buf,
164            lines: Vec::new(),
165            item_count: 0,
166            seen: BTreeMap::new(),
167            path,
168        }
169    }
170}
171
172impl<'a, E: Elem + Into<u32>> Visitor for TreePrinter<'a, E> {
173    fn visit_component(&mut self, name: &str, component: &impl Component) -> core::fmt::Result {
174        match component.ty_name() {
175            "U32Reg" => {
176                // U32regs print as hexadecimal.
177                let mut subtree = TreeGather::new(self.buf);
178                component.walk(&mut subtree)?;
179                if subtree.vals.iter().any(|&x| x != 0) {
180                    let mut msg = format!("{name}: ");
181                    write!(msg, " 0x")?;
182                    for val in subtree.vals {
183                        if val > 255 {
184                            // Unexpected non-byte in u32reg
185                            write!(msg, "[{:#x}?]", val)?;
186                        } else {
187                            write!(msg, "{:02x}", val)?;
188                        }
189                    }
190                    self.lines.push(msg);
191                    self.item_count += 1;
192                }
193            }
194            _ => {
195                // Generic subtree
196                let mut subtree = TreePrinter::new(self.buf, self.path.clone(), name);
197                component.walk(&mut subtree)?;
198                match subtree.item_count {
199                    0 => {}
200                    1 => {
201                        // Single item inside; build up path on a
202                        // single line to save vertical space for
203                        // structures like a{b{c{x: y}}}.
204                        self.item_count += 1;
205                        let mut lines = subtree.lines;
206                        lines[0] = format!("{name}: {{{}", lines[0]);
207                        lines.last_mut().as_mut().unwrap().push_str(" }");
208                        self.lines.extend(lines);
209                    }
210                    _ => {
211                        // Multi-line structure; if we've seen this before, just emit a reference.
212                        // Otherwise, save it.
213                        self.item_count += 1;
214                        match self.seen.get(&subtree.lines) {
215                            Some(path) => {
216                                self.lines.push(format!("{name}: duplicate of {path}"));
217                            }
218                            None => {
219                                self.seen.extend(subtree.seen);
220                                self.lines.push(format!("{name}: {{"));
221                                self.lines
222                                    .extend(subtree.lines.iter().map(|s| format!("  {s}")));
223                                self.lines.push("}".to_string());
224                                self.seen.insert(subtree.lines, subtree.path.join("."));
225                            }
226                        }
227                    }
228                }
229            }
230        }
231        Ok(())
232    }
233
234    fn visit_reg(&mut self, offset: usize) -> core::fmt::Result {
235        if let Some(val) = self.buf.get(offset) {
236            self.lines.push(format!("{:?}", val));
237            self.item_count += 1;
238        }
239        Ok(())
240    }
241}
242
243// Gathers all the registers in a tree as u32s.
244struct TreeGather<'a, E: Elem> {
245    buf: &'a [E],
246    vals: Vec<u32>,
247}
248
249impl<'a, E: Elem> TreeGather<'a, E> {
250    fn new(buf: &'a [E]) -> Self {
251        Self {
252            buf,
253            vals: Vec::new(),
254        }
255    }
256}
257
258impl<'a, E: Elem + Into<u32>> Visitor for TreeGather<'a, E> {
259    fn visit_component(&mut self, _name: &str, component: &impl Component) -> core::fmt::Result {
260        component.walk(self)
261    }
262    fn visit_reg(&mut self, offset: usize) -> core::fmt::Result {
263        if let Some(val) = self.buf.get(offset) {
264            self.vals.push((*val).into())
265        }
266        Ok(())
267    }
268}
risc0_zkp/layout.rs

risc0_zkp/
layout.rs
