plonky2/src/context_tree.rs

use log::debug;

/// The hierarchy of contexts, and the gate count contributed by each one. Useful for debugging.
pub(crate) struct ContextTree {
    /// The name of this scope.
    name: String,
    /// The gate count when this scope was created.
    enter_gate_count: usize,
    /// The gate count when this scope was destroyed, or None if it has not yet been destroyed.
    exit_gate_count: Option<usize>,
    /// Any child contexts.
    children: Vec<ContextTree>,
}

impl ContextTree {
    pub fn new() -> Self {
        Self {
            name: "root".to_string(),
            enter_gate_count: 0,
            exit_gate_count: None,
            children: vec![],
        }
    }

    /// Whether this context is still in scope.
    fn is_open(&self) -> bool {
        self.exit_gate_count.is_none()
    }

    /// A description of the stack of currently-open scopes.
    pub fn open_stack(&self) -> String {
        let mut stack = Vec::new();
        self.open_stack_helper(&mut stack);
        stack.join(" > ")
    }

    fn open_stack_helper(&self, stack: &mut Vec<String>) {
        if self.is_open() {
            stack.push(self.name.clone());
            if let Some(last_child) = self.children.last() {
                last_child.open_stack_helper(stack);
            }
        }
    }

    pub fn push(&mut self, ctx: &str, current_gate_count: usize) {
        assert!(self.is_open());

        if let Some(last_child) = self.children.last_mut() {
            if last_child.is_open() {
                last_child.push(ctx, current_gate_count);
                return;
            }
        }

        self.children.push(ContextTree {
            name: ctx.to_string(),
            enter_gate_count: current_gate_count,
            exit_gate_count: None,
            children: vec![],
        })
    }

    /// Close the deepest open context from this tree.
    pub fn pop(&mut self, current_gate_count: usize) {
        assert!(self.is_open());

        if let Some(last_child) = self.children.last_mut() {
            if last_child.is_open() {
                last_child.pop(current_gate_count);
                return;
            }
        }

        self.exit_gate_count = Some(current_gate_count);
    }

    fn gate_count_delta(&self, current_gate_count: usize) -> usize {
        self.exit_gate_count.unwrap_or(current_gate_count) - self.enter_gate_count
    }

    /// Filter out children with a low gate count.
    pub fn filter(&self, current_gate_count: usize, min_delta: usize) -> Self {
        Self {
            name: self.name.clone(),
            enter_gate_count: self.enter_gate_count,
            exit_gate_count: self.exit_gate_count,
            children: self
                .children
                .iter()
                .filter(|c| c.gate_count_delta(current_gate_count) >= min_delta)
                .map(|c| c.filter(current_gate_count, min_delta))
                .collect(),
        }
    }

    pub fn print(&self, current_gate_count: usize) {
        self.print_helper(current_gate_count, 0);
    }

    fn print_helper(&self, current_gate_count: usize, depth: usize) {
        let prefix = "| ".repeat(depth);
        debug!(
            "{}{} gates to {}",
            prefix,
            self.gate_count_delta(current_gate_count),
            self.name
        );
        for child in &self.children {
            child.print_helper(current_gate_count, depth + 1);
        }
    }
}

/// Creates a named scope; useful for debugging.
#[macro_export]
macro_rules! context {
    ($builder:expr, $ctx:expr, $exp:expr) => {{
        $builder.push_context($ctx);
        let res = $exp;
        $builder.pop_context();
        res
    }};
}