Handle break labels like variables.

jaq-core/src/compile.rs

@@ -1,15 +1,14 @@
 //! Program compilation.
 
 use crate::load::{self, lex, parse};
-use crate::{ops, Bind, Filter};
+use crate::{ops, Bind as Arg, Filter};
 use alloc::collections::{BTreeMap, BTreeSet};
 use alloc::{boxed::Box, string::String, vec::Vec};
 
 type NativeId = usize;
 type ModId = usize;
 type VarId = usize;
 type VarSkip = usize;
-type LabelSkip = usize;
 type Arity = usize;
 
 /// Index of a term in the look-up table.
@@ -81,15 +80,14 @@ pub(crate) enum Term<T = TermId> {
     /// Singleton object (`{f: g}`)
     ObjSingle(T, T),
 
-    /// Bound variable (`$x`) or filter argument (`a`)
-    Var(VarId, LabelSkip),
+    /// Bound variable (`$x`), label (`label $x`), or filter argument (`a`)
+    Var(VarId),
     /// Call to a filter (`filter`, `filter(…)`)
-    CallDef(TermId, Box<[Bind<T>]>, VarSkip, Option<Tailrec>),
-    Native(NativeId, Box<[Bind<T>]>),
+    CallDef(TermId, Box<[Arg<T>]>, VarSkip, Option<Tailrec>),
+    Native(NativeId, Box<[Arg<T>]>),
 
+    /// Binding of a break label (`label $x | f`)
     Label(T),
-    Break(usize),
-
     /// Negation operation (`-f`)
     Neg(T),
     /// Variable binding (`f as $x | g`) if identifier (`x`) is given, otherwise
@@ -233,22 +231,22 @@ impl<S: Default, F> Default for Compiler<S, F> {
 }
 
 #[derive(Clone, Debug)]
-struct Sig<S, A = Bind> {
+struct Sig<S, A = Arg> {
     name: S,
     // TODO: we could analyse for each argument whether it is TR, and
     // use this when converting args at callsite
     args: Box<[A]>,
 }
 
-fn bind_from<T>(s: &str, x: T) -> Bind<T> {
+fn bind_from<T>(s: &str, x: T) -> Arg<T> {
     if s.starts_with('$') {
-        Bind::Var(x)
+        Arg::Var(x)
     } else {
-        Bind::Fun(x)
+        Arg::Fun(x)
     }
 }
 
-fn binds<T, U: Copy>(binds: &[Bind<T>], args: &[U]) -> Box<[Bind<U>]> {
+fn binds<T, U: Copy>(binds: &[Arg<T>], args: &[U]) -> Box<[Arg<U>]> {
     assert!(binds.len() == args.len());
     let args = binds.iter().zip(args);
     args.map(|(bind, id)| bind.as_ref().map(|_| *id)).collect()
@@ -270,7 +268,7 @@ impl<S: Eq, A> Sig<S, A> {
 }
 
 impl Def {
-    fn call(&self, args: Box<[Bind<TermId>]>, vars: usize) -> Term {
+    fn call(&self, args: Box<[Arg<TermId>]>, vars: usize) -> Term {
         // we pretend that the function call is tail-recursive,
         // and at the very end of compilation, we will correct calls
         // to non-tail-recursive functions
@@ -279,12 +277,20 @@ impl Def {
 }
 
 /// Store a map of vectors plus the sum of the lengths of all vectors.
-#[derive(Default)]
 struct MapVecLen<S> {
     bound: MapVec<S, usize>,
     total: usize,
 }
 
+impl<S> Default for MapVecLen<S> {
+    fn default() -> Self {
+        Self {
+            bound: MapVec::default(),
+            total: 0,
+        }
+    }
+}
+
 impl<S: Ord> MapVecLen<S> {
     fn push(&mut self, name: S) {
         self.total += 1;
@@ -302,22 +308,40 @@ impl<S: Ord> MapVecLen<S> {
 }
 
 enum Fun<S> {
-    // number of labels
-    Arg(usize),
-    Parent(Box<[Bind<S>]>, Def),
+    Arg,
+    Parent(Box<[Arg<S>]>, Def),
     // Tr is for tail-rec allowed funs
-    Sibling(Box<[Bind<S>]>, Def, Tr),
+    Sibling(Box<[Arg<S>]>, Def, Tr),
+}
+
+/// Single binding.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub(crate) enum Bind<V, L = V, F = V> {
+    /// binding to a variable
+    Var(V),
+    /// binding to a break label
+    Label(L),
+    /// binding to a filter
+    Fun(F),
 }
 
-#[derive(Default)]
 struct Locals<S> {
     // usize = number of vars
     funs: MapVec<(S, Arity), (Fun<S>, usize)>,
-    vars: MapVecLen<S>,
-    labels: MapVecLen<S>,
+    vars: MapVecLen<Bind<S>>,
     parents: Tr,
 }
 
+impl<S> Default for Locals<S> {
+    fn default() -> Self {
+        Self {
+            funs: MapVec::default(),
+            vars: MapVecLen::default(),
+            parents: Tr::default(),
+        }
+    }
+}
+
 struct MapVec<K, V>(BTreeMap<K, Vec<V>>);
 
 impl<K, V> Default for MapVec<K, V> {
@@ -354,15 +378,15 @@ impl<K: Ord, V> MapVec<K, V> {
 }
 
 impl<S: Copy + Ord> Locals<S> {
-    fn push_sibling(&mut self, name: S, args: Box<[Bind<S>]>, def: Def) {
+    fn push_sibling(&mut self, name: S, args: Box<[Arg<S>]>, def: Def) {
         let tr = self.parents.clone();
         self.funs.push(
             (name, args.len()),
             (Fun::Sibling(args, def, tr), self.vars.total),
         );
     }
 
-    fn pop_sibling(&mut self, name: S, arity: Arity) -> (Box<[Bind<S>]>, Def, Tr) {
+    fn pop_sibling(&mut self, name: S, arity: Arity) -> (Box<[Arg<S>]>, Def, Tr) {
         let (y, vars) = match self.funs.pop(&(name, arity)) {
             Some((Fun::Sibling(args, def, tr), vars)) => ((args, def, tr), vars),
             _ => panic!(),
@@ -371,15 +395,15 @@ impl<S: Copy + Ord> Locals<S> {
         y
     }
 
-    fn push_parent(&mut self, name: S, args: Box<[Bind<S>]>, def: Def) {
+    fn push_parent(&mut self, name: S, args: Box<[Arg<S>]>, def: Def) {
         self.parents.insert(def.id);
 
         let vars = self.vars.total;
 
         for arg in args.iter() {
             match arg {
-                Bind::Var(v) => self.vars.push(*v),
-                Bind::Fun(f) => self.push_arg(*f),
+                Arg::Var(v) => self.vars.push(Bind::Var(*v)),
+                Arg::Fun(f) => self.push_arg(*f),
             }
         }
 
@@ -394,8 +418,8 @@ impl<S: Copy + Ord> Locals<S> {
         };
         for arg in args.iter().rev() {
             match arg {
-                Bind::Var(v) => self.vars.pop(v),
-                Bind::Fun(f) => self.pop_arg(*f),
+                Arg::Var(v) => self.vars.pop(&Bind::Var(*v)),
+                Arg::Fun(f) => self.pop_arg(*f),
             }
         }
         assert_eq!(self.vars.total, vars);
@@ -404,26 +428,22 @@ impl<S: Copy + Ord> Locals<S> {
     }
 
     fn push_arg(&mut self, name: S) {
-        self.vars.total += 1;
-        self.funs
-            .push((name, 0), (Fun::Arg(self.labels.total), self.vars.total));
+        self.vars.push(Bind::Fun(name));
+        self.funs.push((name, 0), (Fun::Arg, self.vars.total));
     }
 
     fn pop_arg(&mut self, name: S) {
-        let (labels, vars) = match self.funs.pop(&(name, 0)) {
-            Some((Fun::Arg(labels), vars)) => (labels, vars),
+        let vars = match self.funs.pop(&(name, 0)) {
+            Some((Fun::Arg, vars)) => vars,
             _ => panic!(),
         };
-        assert_eq!(self.labels.total, labels);
         assert_eq!(self.vars.total, vars);
-        self.vars.total -= 1;
+        self.vars.pop(&Bind::Fun(name));
     }
 
     fn call(&mut self, name: S, args: &[TermId], tr: &Tr) -> Option<Term> {
         Some(match self.funs.get_last_mut(&(name, args.len()))? {
-            (Fun::Arg(labels), vars) => {
-                Term::Var(self.vars.total - *vars, self.labels.total - *labels)
-            }
+            (Fun::Arg, vars) => Term::Var(self.vars.total - *vars),
             (Fun::Sibling(args_, def, tr_), vars) => {
                 // we are at a position that may only call `tr` tail-recursively and
                 // we call a sibling that may only call `tr_` tail-recursively,
@@ -447,10 +467,7 @@ impl<S: Copy + Ord> Locals<S> {
     }
 
     fn is_empty(&self) -> bool {
-        self.funs.is_empty()
-            && self.vars.is_empty()
-            && self.labels.is_empty()
-            && self.parents.is_empty()
+        self.funs.is_empty() && self.vars.is_empty() && self.parents.is_empty()
     }
 }
 
@@ -460,7 +477,7 @@ type Tr = BTreeSet<TermId>;
 
 impl<'s, F> Compiler<&'s str, F> {
     /// Supply functions with given signatures.
-    pub fn with_funs(mut self, funs: impl IntoIterator<Item = (&'s str, Box<[Bind]>, F)>) -> Self {
+    pub fn with_funs(mut self, funs: impl IntoIterator<Item = (&'s str, Box<[Arg]>, F)>) -> Self {
         self.lut.funs = funs
             .into_iter()
             .map(|(name, args, f)| (Sig { name, args }, f))
@@ -533,16 +550,19 @@ impl<'s, F> Compiler<&'s str, F> {
     }
 
     fn with_label<T>(&mut self, label: &'s str, f: impl FnOnce(&mut Self) -> T) -> T {
-        self.locals.labels.push(label);
+        self.locals.vars.push(Bind::Label(label));
         let y = f(self);
-        self.locals.labels.pop(&label);
+        self.locals.vars.pop(&Bind::Label(label));
         y
     }
 
     fn with_vars<T>(&mut self, vars: &[&'s str], f: impl FnOnce(&mut Self) -> T) -> T {
-        vars.iter().for_each(|v| self.locals.vars.push(v));
+        vars.iter()
+            .for_each(|v| self.locals.vars.push(Bind::Var(v)));
         let y = f(self);
-        vars.iter().rev().for_each(|v| self.locals.vars.pop(v));
+        vars.iter()
+            .rev()
+            .for_each(|v| self.locals.vars.pop(&Bind::Var(v)));
         y
     }
 
@@ -583,7 +603,7 @@ impl<'s, F> Compiler<&'s str, F> {
     }
 
     /// Compile a placeholder sibling with its corresponding definition.
-    fn def_post(&mut self, d: parse::Def<&'s str>) -> (Sig<&'s str, Bind>, Def) {
+    fn def_post(&mut self, d: parse::Def<&'s str>) -> (Sig<&'s str, Arg>, Def) {
         let (args, def, mut tr) = self.locals.pop_sibling(d.name, d.args.len());
         let tid = def.id;
         self.locals.push_parent(d.name, args, def);
@@ -808,19 +828,19 @@ impl<'s, F> Compiler<&'s str, F> {
     fn var(&mut self, x: &'s str) -> Term {
         let mut i = self.locals.vars.total;
 
-        if let Some(v) = self.locals.vars.bound.get_last(&x) {
-            return Term::Var(i - v, 0);
+        if let Some(v) = self.locals.vars.bound.get_last(&Bind::Var(x)) {
+            return Term::Var(i - v);
         }
         for (x_, mid) in self.imported_vars.iter().rev() {
             if x == *x_ && *mid == self.mod_map.len() {
-                return Term::Var(i, 0);
+                return Term::Var(i);
             } else {
                 i += 1;
             }
         }
         for x_ in self.global_vars.iter().rev() {
             if x == *x_ {
-                return Term::Var(i, 0);
+                return Term::Var(i);
             } else {
                 i += 1;
             }
@@ -829,8 +849,8 @@ impl<'s, F> Compiler<&'s str, F> {
     }
 
     fn break_(&mut self, x: &'s str) -> Term {
-        if let Some(l) = self.locals.labels.bound.get_last(&x) {
-            return Term::Break(self.locals.labels.total - l);
+        if let Some(l) = self.locals.vars.bound.get_last(&Bind::Label(x)) {
+            return Term::Var(self.locals.vars.total - l);
         }
         self.fail(x, Undefined::Label)
     }

jaq-core/src/exn.rs

@@ -17,7 +17,7 @@ pub(crate) enum Inner<'a, V> {
     ///
     /// This is used internally to execute tail-recursive filters.
     /// If this can be observed by users, then this is a bug.
-    TailCall(&'a crate::compile::TermId, crate::Vars<'a, V>, V),
+    TailCall(&'a crate::compile::TermId, crate::filter::Vars<'a, V>, V),
     Break(usize),
 }