From 728c71f96dca81a7324bdd98bd54d1d450029802 Mon Sep 17 00:00:00 2001
From: Bruno Andreotti <bandreotti@proton.me>
Date: Thu, 16 May 2024 14:54:20 -0300
Subject: [PATCH] Recompute outbound premises when elaborating

---
 carcara/src/ast/node.rs       | 13 +++++++++++++
 carcara/src/elaborator/mod.rs | 26 +++++++++++++++++++++-----
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/carcara/src/ast/node.rs b/carcara/src/ast/node.rs
index af599160..c8f974bc 100644
--- a/carcara/src/ast/node.rs
+++ b/carcara/src/ast/node.rs
@@ -60,6 +60,19 @@ impl ProofNode {
         }
     }
 
+    /// Returns a vector of the "outbound" premises of this node.
+    ///
+    /// These are the premises whose depth is smaller than the node's depth, that is, the premises
+    /// that refer to outside of this node's subproof.
+    pub fn get_outbound_premises(&self) -> Vec<Rc<ProofNode>> {
+        let ps = match self {
+            ProofNode::Assume { .. } => return Vec::new(),
+            ProofNode::Step(s) => s.premises.iter(),
+            ProofNode::Subproof(s) => s.outbound_premises.iter(),
+        };
+        ps.filter(|p| p.depth() < self.depth()).cloned().collect()
+    }
+
     /// Returns `true` if the node is an `assume` command.
     pub fn is_assume(&self) -> bool {
         matches!(self, ProofNode::Assume { .. })
diff --git a/carcara/src/elaborator/mod.rs b/carcara/src/elaborator/mod.rs
index 6e921eb6..b5122333 100644
--- a/carcara/src/elaborator/mod.rs
+++ b/carcara/src/elaborator/mod.rs
@@ -1,6 +1,7 @@
 mod transitivity;
 
 use crate::ast::*;
+use indexmap::IndexSet;
 use std::collections::{HashMap, HashSet};
 
 #[allow(unused)]
@@ -22,9 +23,12 @@ where
     let mut cache: HashMap<&Rc<ProofNode>, Rc<ProofNode>> = HashMap::new();
     let mut did_outbound: HashSet<&Rc<ProofNode>> = HashSet::new();
     let mut todo = vec![(root, false)];
+    let mut outbound_premises_stack = vec![IndexSet::new()];
 
     while let Some((node, is_done)) = todo.pop() {
-        assert!(!cache.contains_key(node), "this shouldn't happen I think");
+        if cache.contains_key(node) {
+            continue;
+        }
 
         let mutated = match node.as_ref() {
             ProofNode::Assume { .. } => mutate_func(pool, node),
@@ -57,6 +61,11 @@ where
                 mutate_func(pool, &new_node)
             }
             ProofNode::Subproof(s) if !is_done => {
+                assert!(
+                    node.depth() == outbound_premises_stack.len() - 1,
+                    "all outbound premises should have already been dealt with!"
+                );
+
                 if !did_outbound.contains(node) {
                     did_outbound.insert(node);
                     todo.push((node, false));
@@ -66,18 +75,25 @@ where
 
                 todo.push((node, true));
                 todo.push((&s.last_step, false));
+                outbound_premises_stack.push(IndexSet::new());
                 continue;
             }
             ProofNode::Subproof(s) => {
-                let new_node = Rc::new(ProofNode::Subproof(SubproofNode {
+                let outbound_premises =
+                    outbound_premises_stack.pop().unwrap().into_iter().collect();
+                Rc::new(ProofNode::Subproof(SubproofNode {
                     last_step: cache[&s.last_step].clone(),
                     args: s.args.clone(),
-                    outbound_premises: Vec::new(), // TODO: recompute outbound premises
-                }));
-                mutate_func(pool, &new_node)
+                    outbound_premises,
+                }))
             }
         };
+        outbound_premises_stack
+            .last_mut()
+            .unwrap()
+            .extend(mutated.get_outbound_premises());
         cache.insert(node, mutated);
     }
+    assert!(outbound_premises_stack.len() == 1 && outbound_premises_stack[0].is_empty());
     cache[root].clone()
 }