Add symbol collection visitor

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>

src/compile/ctx.rs

@@ -1,13 +1,19 @@
-use crate::compile::name::NameStack;
+use crate::{
+    compile::name::NameStack,
+    obj::Sym,
+};
+use std::collections::HashMap;
 
 pub struct Ctx {
     name_stack: NameStack,
+    syms: HashMap<String, Sym>,
 }
 
 impl Ctx {
     pub fn new() -> Self {
         Ctx {
             name_stack: NameStack::new(),
+            syms: Default::default(),
         }
     }
 
@@ -18,4 +24,19 @@ impl Ctx {
     pub fn name_stack_mut(&mut self) -> &mut NameStack {
         &mut self.name_stack
     }
+
+    pub fn syms(&self) -> &HashMap<String, Sym> {
+        &self.syms
+    }
+
+    pub fn syms_mut(&mut self) -> &mut HashMap<String, Sym> {
+        &mut self.syms
+    }
+
+    pub fn add_sym(&mut self, name: String) -> Sym {
+        let next_sym = self.syms().len();
+        *self.syms_mut()
+            .entry(name)
+            .or_insert(next_sym.into())
+    }
 }

src/compile/name.rs

@@ -9,18 +9,25 @@ use std::collections::HashMap;
 // NameStack
 ////////////////////////////////////////////////////////////////////////////////
 
+#[derive(Debug, Clone)]
 pub struct NameStack {
     next_sym: usize,
     stack: Vec<HashMap<String, NameId>>,
 }
 
-impl NameStack {
-    pub fn new() -> Self {
+impl Default for NameStack {
+    fn default() -> Self {
         NameStack {
             next_sym: 0,
             stack: vec![Default::default()],
         }
     }
+}
+
+impl NameStack {
+    pub fn new() -> Self {
+        Default::default()
+    }
 
     pub fn push(&mut self, syms: HashMap<String, NameId>) {
         self.stack.push(syms);
@@ -155,3 +162,110 @@ impl Visit<BaseExpr> for CollectNames<'_, '_> {
 ////////////////////////////////////////////////////////////////////////////////
 // CollectSyms
 ////////////////////////////////////////////////////////////////////////////////
+
+/// Collect symbols (identifiers and symbols included).
+///
+/// These are used for indexing values in attribute maps.
+pub struct CollectSyms<'c, 't> {
+    ctx: &'c mut Ctx,
+    text: &'t str,
+}
+
+impl<'c, 't> CollectSyms<'c, 't> {
+    pub fn new(ctx: &'c mut Ctx, text: &'t str) -> Self {
+        CollectSyms { ctx, text }
+    }
+
+    pub fn collect(&mut self, stmts: &Vec<Stmt>) {
+        // Collect all LHS assignments
+        for stmt in stmts.iter() {
+            self.visit(stmt);
+        }
+    }
+}
+
+default_visitor!(Stmt for CollectSyms<'_, '_> where Out = ());
+default_visitor!(Expr for CollectSyms<'_, '_> where Out = ());
+
+impl Visit<AssignStmt> for CollectSyms<'_, '_> {
+    type Out = ();
+
+    fn visit(&mut self, stmt: &AssignStmt) -> Self::Out {
+        self.visit(&stmt.lhs);
+        self.visit(&stmt.rhs);
+    }
+}
+
+impl Visit<BinExpr> for CollectSyms<'_, '_> {
+    type Out = ();
+
+    fn visit(&mut self, expr: &BinExpr) -> Self::Out {
+        self.visit(&expr.lhs);
+        self.visit(&expr.rhs);
+    }
+}
+
+impl Visit<UnExpr> for CollectSyms<'_, '_> {
+    type Out = ();
+
+    fn visit(&mut self, expr: &UnExpr) -> Self::Out {
+        self.visit(&expr.expr);
+    }
+}
+
+impl Visit<FunCallExpr> for CollectSyms<'_, '_> {
+    type Out = ();
+
+    fn visit(&mut self, expr: &FunCallExpr) -> Self::Out {
+        self.visit(&expr.expr);
+        expr.args.iter().for_each(|expr| self.visit(expr));
+    }
+}
+
+impl Visit<IndexExpr> for CollectSyms<'_, '_> {
+    type Out = ();
+
+    fn visit(&mut self, expr: &IndexExpr) -> Self::Out {
+        self.visit(&expr.expr);
+        self.visit(&expr.index);
+    }
+}
+
+impl Visit<FunExpr> for CollectSyms<'_, '_> {
+    type Out = ();
+
+    fn visit(&mut self, expr: &FunExpr) -> Self::Out {
+        expr.params.iter().for_each(|name| {
+            self.ctx.add_sym(name.clone());
+        });
+        self.visit(&expr.expr);
+    }
+}
+
+impl Visit<BaseExpr> for CollectSyms<'_, '_> {
+    type Out = ();
+
+    fn visit(&mut self, expr: &BaseExpr) -> Self::Out {
+        match &expr.kind {
+            BaseExprKind::Ident => {
+                let name = expr.text_at(self.text).to_string();
+                self.ctx.add_sym(name);
+            }
+            BaseExprKind::Sym => {
+                let name = expr.text_at(self.text).chars().skip(1).collect::<String>();
+                self.ctx.add_sym(name);
+            }
+            BaseExprKind::List(l) | BaseExprKind::Tuple(l) => {
+                l.iter().for_each(|expr| self.visit(expr))
+            }
+            BaseExprKind::Object(o) => o.iter().for_each(|(k, v)| {
+                self.visit(k);
+                self.visit(v);
+            }),
+            BaseExprKind::Block(stmts) => {
+                stmts.iter().for_each(|stmt| self.visit(stmt));
+            }
+            _ => { /* no-op */ }
+        }
+    }
+}

src/main.rs

@@ -44,13 +44,16 @@ fn main() -> Result<()> {
     //println!("{:#?}", ast);
     
     let mut ctx = compile::ctx::Ctx::new();
+    compile::name::CollectSyms::new(&mut ctx, text.as_str()).collect(&ast);
+    println!("{:#?}", ctx.syms());
+
     ctx.name_stack_mut().push_default();
     {
         let mut collect_names = compile::name::CollectNames::new(&mut ctx, text.as_str());
         collect_names.collect(&ast);
     }
     let names = ctx.name_stack_mut().pop().unwrap();
-    println!("{:#?}", names);
+    //println!("{:#?}", names);
     
     Ok(())
 }