finish 1b

examples/1a.aoc

@@ -3,8 +3,8 @@ input data = "1.txt";
 let lines = data.splitlines();
 
 let pairs = lines.map((s) => {
-        map(s.splitwhitespace(), (n) => { parseint(n) })
-    });
+    map(s.splitwhitespace(), (n) => { parseint(n) })
+});
 
 let lists = pairs.transpose();
 
@@ -18,10 +18,10 @@ let lists = [left, right];
 let lists = lists.transpose();
 
 let differences = lists.map((a) => { 
-        let left = a[0];
-        let right = a[1]; 
-        abs(left - right)
-    });
+    let left = a[0];
+    let right = a[1]; 
+    abs(left - right)
+});
 print(differences);
 
 let result = differences.sum();

examples/1b.aoc

@@ -0,0 +1,24 @@
+input data = "1.txt";
+
+let lines = data.splitlines();
+
+let pairs = lines.map((s) => {
+    map(s.splitwhitespace(), (n) => { parseint(n) })
+});
+
+let lists = pairs.transpose();
+
+let left = lists[0];
+let right = lists[1];
+
+let counts = left.map((n) => {
+    let one_count = right.map((m) => {
+        if m == n then 1 else 0
+    });
+    let s = one_count.sum();
+    n * s
+});
+
+let result = counts.sum();
+
+print(result);

src/ast.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use anyhow::Result;
 use derive_more::Debug;
 
@@ -12,6 +14,8 @@ pub enum Stmt {
 pub enum Expr {
   Block(Vec<Stmt>, Option<Box<Expr>>),
   Array(Vec<Expr>),
+  Bool(bool),
+  IfThenElse(Box<Expr>, Box<Expr>, Box<Expr>),
   Int(i64),
   Ident(String),
   BinOp(Box<Expr>, Op, Box<Expr>),
@@ -24,14 +28,17 @@ pub enum Expr {
 #[derive(Debug, Clone)]
 pub enum Op {
   Sub,
+  Mul,
+  EqEq,
 }
 
 #[derive(Debug, Clone)]
 pub enum Value {
   Unit,
+  Bool(bool),
   Int(i64),
   Array(Vec<Value>),
   String(String),
-  Function(Vec<String>, Expr),
+  Function(HashMap<String, Value>, Vec<String>, Expr),
   Builtin(#[debug(skip)] fn(Vec<Value>) -> Result<Value>),
 }

src/grammar.lalrpop

@@ -21,11 +21,18 @@ Expr: Expr = {
   
   #[precedence(level = "1")] #[assoc(side = "left")]
   <l:Expr> "-" <r:Expr> => Expr::BinOp(Box::new(l), Op::Sub, Box::new(r)),
+  
+  #[precedence(level = "2")] #[assoc(side = "left")]
+  <l:Expr> "*" <r:Expr> => Expr::BinOp(Box::new(l), Op::Mul, Box::new(r)),
+  <l:Expr> "==" <r:Expr> => Expr::BinOp(Box::new(l), Op::EqEq, Box::new(r)),
 
-  #[precedence(level = "2")]
+  #[precedence(level = "3")]
+  "if" <c:Expr> "then" <l:Expr> "else" <r:Expr> => Expr::IfThenElse(Box::new(c), Box::new(l), Box::new(r)),
+  
+  #[precedence(level = "4")]
   <e:Expr> "." <n:Ident> => Expr::Field(Box::new(e), n),
   
-  #[precedence(level = "3")]
+  #[precedence(level = "5")]
   <e:Expr> "(" <a:Sep<",", Expr>> ")" => Expr::Call(Box::new(e), a),
   <e:Expr> "[" <i:Expr> "]" => Expr::Index(Box::new(e), Box::new(i)),
   

src/main.rs

@@ -1,8 +1,14 @@
 mod ast;
 
 use std::{
-  collections::HashMap, env, fmt::format, fs, hash::Hash, os::unix::thread,
-  path::PathBuf, process::exit,
+  collections::{HashMap, HashSet},
+  env,
+  fmt::format,
+  fs,
+  hash::Hash,
+  os::unix::thread,
+  path::PathBuf,
+  process::exit,
 };
 
 use anyhow::{bail, ensure, Context, Result};
@@ -105,7 +111,7 @@ fn eval_expr(ctx: &mut HashMap<String, Value>, expr: &Expr) -> Result<Value> {
   match expr {
     Expr::Ident(n) => match ctx.get(n) {
       Some(v) => Ok(v.clone()),
-      None => bail!("unknown name {n}"),
+      None => bail!("{}:{}: unknown name {n}", file!(), line!()),
     },
     Expr::Call(expr, args) if matches!(**expr, Expr::Field(_, _)) => {
       if let Expr::Field(expr, name) = &**expr {
@@ -124,7 +130,7 @@ fn eval_expr(ctx: &mut HashMap<String, Value>, expr: &Expr) -> Result<Value> {
     Expr::Call(expr, args) => {
       let v = eval_expr(ctx, expr)?;
       match v {
-        Value::Function(params, body) => {
+        Value::Function(env, params, body) => {
           let mapping = params.iter().zip(args.iter()).collect::<Vec<_>>();
           todo!()
         }
@@ -139,8 +145,21 @@ fn eval_expr(ctx: &mut HashMap<String, Value>, expr: &Expr) -> Result<Value> {
       }
     }
     Expr::Lambda(vec, expr) => {
-      // TODO: Do free variable analysis
-      Ok(Value::Function(vec.clone(), (**expr).clone()))
+      let mut free_vars = HashSet::new();
+      let mut known_vars =
+        vec.iter().map(|s| s.to_owned()).collect::<HashSet<_>>();
+      free_variables_expr_rec(expr, &mut known_vars, &mut free_vars);
+
+      let mut env = HashMap::new();
+      for var in free_vars {
+        let v = match ctx.get(&var) {
+          Some(v) => v,
+          None => bail!("{}:{}: unknown name {var}", file!(), line!()),
+        };
+        env.insert(var, v.clone());
+      }
+
+      Ok(Value::Function(env, vec.clone(), (**expr).clone()))
     }
     Expr::Block(vec, expr) => {
       for stmt in vec {
@@ -182,6 +201,29 @@ fn eval_expr(ctx: &mut HashMap<String, Value>, expr: &Expr) -> Result<Value> {
           (Value::Int(left), Value::Int(right)) => Ok(Value::Int(left - right)),
           _ => bail!("both {left:?} , {right:?} must be ints"),
         },
+        Op::Mul => match (&left, &right) {
+          (Value::Int(left), Value::Int(right)) => Ok(Value::Int(left * right)),
+          _ => bail!("both {left:?} , {right:?} must be ints"),
+        },
+        Op::EqEq => match (&left, &right) {
+          (Value::Int(left), Value::Int(right)) => {
+            Ok(Value::Bool(left == right))
+          }
+          _ => bail!("both {left:?} , {right:?} must be the same"),
+        },
+      }
+    }
+    Expr::Bool(b) => Ok(Value::Bool(*b)),
+    Expr::IfThenElse(cond, t, f) => {
+      let cond = match eval_expr(ctx, cond)? {
+        Value::Bool(b) => b,
+        _ => bail!("not a bool"),
+      };
+
+      if cond {
+        Ok(eval_expr(ctx, t)?)
+      } else {
+        Ok(eval_expr(ctx, f)?)
       }
     }
   }
@@ -222,10 +264,14 @@ fn insert_builtins(ctx: &mut HashMap<String, Value>) {
     let list = &args[0];
     let func = &args[1];
     match (list, func) {
-      (Value::Array(a), Value::Function(s, f)) => {
+      (Value::Array(a), Value::Function(env, s, f)) => {
         let mut res = Vec::new();
         let mut ctx = HashMap::new();
         insert_builtins(&mut ctx);
+        for (k, v) in env {
+          ctx.insert(k.to_owned(), v.clone());
+        }
+
         for val in a {
           ctx.insert(s[0].to_owned(), val.clone());
           res.push(eval_expr(&mut ctx, f)?);
@@ -334,3 +380,76 @@ fn insert_builtins(ctx: &mut HashMap<String, Value>) {
   }
   ctx.insert("sum".to_owned(), Value::Builtin(sum));
 }
+
+fn free_variables_stmt_rec(
+  st: &Stmt,
+  ctx: &mut HashSet<String>,
+  s: &mut HashSet<String>,
+) {
+  match st {
+    Stmt::Input(_, _) => {}
+    Stmt::Let(n, expr) => {
+      free_variables_expr_rec(expr, ctx, s);
+      ctx.insert(n.to_owned());
+    }
+    Stmt::Print(expr) => free_variables_expr_rec(expr, ctx, s),
+  }
+}
+
+fn free_variables_expr_rec(
+  e: &Expr,
+  ctx: &mut HashSet<String>,
+  s: &mut HashSet<String>,
+) {
+  match e {
+    Expr::Block(vec, expr) => {
+      for v in vec {
+        free_variables_stmt_rec(v, ctx, s);
+      }
+      if let Some(expr) = expr {
+        free_variables_expr_rec(expr, ctx, s);
+      }
+    }
+    Expr::Array(vec) => {
+      for v in vec {
+        free_variables_expr_rec(v, ctx, s);
+      }
+    }
+    Expr::Bool(_) => {}
+    Expr::Int(_) => {}
+    Expr::Ident(n) => {
+      if !ctx.contains(n) {
+        s.insert(n.to_owned());
+      }
+    }
+    Expr::IfThenElse(expr, expr1, expr2) => {
+      free_variables_expr_rec(expr, ctx, s);
+      free_variables_expr_rec(expr1, ctx, s);
+      free_variables_expr_rec(expr2, ctx, s);
+    }
+    Expr::BinOp(expr, _, expr1) => {
+      free_variables_expr_rec(expr, ctx, s);
+      free_variables_expr_rec(expr1, ctx, s);
+    }
+    Expr::Field(expr, _) => {
+      free_variables_expr_rec(expr, ctx, s);
+    }
+    Expr::Call(expr, vec) => {
+      free_variables_expr_rec(expr, ctx, s);
+      for v in vec {
+        free_variables_expr_rec(v, ctx, s);
+      }
+    }
+    Expr::Index(expr, expr1) => {
+      free_variables_expr_rec(expr, ctx, s);
+      free_variables_expr_rec(expr1, ctx, s);
+    }
+    Expr::Lambda(vec, expr) => {
+      let mut new_set = ctx.clone();
+      for v in vec {
+        new_set.insert(v.to_owned());
+      }
+      free_variables_expr_rec(expr, &mut new_set, s);
+    }
+  }
+}