Impl Drop for ast::Tree to avoid stack overflows

HigherOrderCO · Apr 5, 2024 · a0ea0db · a0ea0db
1 parent 503a548
commit a0ea0db
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 12 deletions.
diff --git a/src/ast.rs b/src/ast.rs
@@ -184,8 +184,12 @@ impl Tree {
     }
   }
 
-  pub fn legacy_mat(arms: Tree, out: Tree) -> Option<Tree> {
-    let Tree::Ctr { lab: 0, ports } = arms else { None? };
+  pub fn legacy_mat(mut arms: Tree, out: Tree) -> Option<Tree> {
+    let ports = if let Tree::Ctr { lab: 0, ports } = &mut arms {
+      std::mem::take(ports)
+    } else {
+      return None;
+    };
     let Ok([zero, succ]) = <[_; 2]>::try_from(ports) else { None? };
     let zero = Box::new(zero);
     let succ = Box::new(succ);
@@ -529,6 +533,7 @@ impl fmt::Display for Tree {
     })
   }
 }
+
 impl Clone for Tree {
   fn clone(&self) -> Tree {
     maybe_grow(|| match self {
@@ -548,3 +553,36 @@ impl Clone for Tree {
     })
   }
 }
+
+impl Drop for Tree {
+  fn drop(&mut self) {
+    fn take_children(tree: &mut Tree, stack: &mut Vec<Tree>) {
+      match tree {
+        Tree::Ctr { ports, .. } => stack.append(ports),
+        Tree::Op { rhs, out, .. } => {
+          stack.push(std::mem::take(rhs.as_mut()));
+          stack.push(std::mem::take(out.as_mut()));
+        }
+        Tree::Mat { zero, succ, out } => {
+          stack.push(std::mem::take(zero.as_mut()));
+          stack.push(std::mem::take(succ.as_mut()));
+          stack.push(std::mem::take(out.as_mut()));
+        }
+        Tree::Adt { fields, .. } => stack.append(fields),
+        Tree::Era | Tree::Num { .. } | Tree::Ref { .. } | Tree::Var { .. } => {}
+      }
+    }
+
+    // Shortcut if we know it has no children
+    if matches!(self, Tree::Era | Tree::Num { .. } | Tree::Ref { .. } | Tree::Var { .. }) {
+      return;
+    }
+
+    let mut stack = vec![];
+    take_children(self, &mut stack);
+
+    while let Some(mut term) = stack.pop() {
+      take_children(&mut term, &mut stack);
+    }
+  }
+}
diff --git a/src/host/encode.rs b/src/host/encode.rs
@@ -61,13 +61,14 @@ impl<'a, E: Encoder> State<'a, E> {
     self.visit_tree(tree, trg);
   }
   fn visit_tree(&mut self, tree: &'a Tree, trg: E::Trg) {
+    static ERA: Tree = Tree::Era;
     maybe_grow(move || match tree {
       Tree::Era => self.encoder.link_const(trg, Port::ERA),
       Tree::Num { val } => self.encoder.link_const(trg, Port::new_num(*val)),
       Tree::Ref { nam } => self.encoder.link_const(trg, Port::new_ref(&self.host.defs[nam])),
       Tree::Ctr { lab, ports } => {
         if ports.is_empty() {
-          return self.visit_tree(&Tree::Era, trg);
+          return self.visit_tree(&tree, trg);
         }
         let mut trg = trg;
         for port in &ports[0 .. ports.len() - 1] {
@@ -81,7 +82,7 @@ impl<'a, E: Encoder> State<'a, E> {
         let mut trg = trg;
         for _ in 0 .. *variant_index {
           let (l, r) = self.encoder.ctr(*lab, trg);
-          self.visit_tree(&Tree::Era, l);
+          self.visit_tree(&ERA, l);
           trg = r;
         }
         let (mut l, mut r) = self.encoder.ctr(*lab, trg);
@@ -92,7 +93,7 @@ impl<'a, E: Encoder> State<'a, E> {
         }
         for _ in 0 .. (*variant_count - *variant_index - 1) {
           let (x, y) = self.encoder.ctr(*lab, r);
-          self.visit_tree(&Tree::Era, x);
+          self.visit_tree(&ERA, x);
           r = y;
         }
         self.encoder.link(l, r);

diff --git a/src/transform/coalesce_ctrs.rs b/src/transform/coalesce_ctrs.rs
@@ -10,13 +10,13 @@ impl Tree {
     maybe_grow(|| match self {
       Tree::Ctr { lab, ports } => {
         ports.iter_mut().for_each(Tree::coalesce_constructors);
-        match ports.pop() {
-          Some(Tree::Ctr { lab: inner_lab, ports: mut inner_ports })
-            if inner_lab == *lab && ports.len() + inner_ports.len() < MAX_ARITY =>
+        match &mut ports.pop() {
+          Some(Tree::Ctr { lab: inner_lab, ports: inner_ports })
+            if inner_lab == lab && ports.len() + inner_ports.len() < MAX_ARITY =>
           {
             ports.extend(inner_ports.drain(..));
           }
-          Some(other) => ports.push(other),
+          Some(other) => ports.push(std::mem::take(other)),
           None => (),
         }
       }

diff --git a/src/transform/encode_adts.rs b/src/transform/encode_adts.rs
@@ -35,10 +35,10 @@ impl Tree {
         }
 
         if let Some((start, idx)) = get_adt_info(lab, ports) {
-          let fields = match ports.swap_remove(idx) {
-            Tree::Ctr { ports: mut fields, .. } => {
+          let fields = match &mut ports.swap_remove(idx) {
+            Tree::Ctr { ports: fields, .. } => {
               fields.pop();
-              fields
+              std::mem::take(fields)
             }
             Tree::Var { .. } => vec![],
             _ => unreachable!(),