From 0f826829f43ff294531a2cc8bc5289d34f6510e2 Mon Sep 17 00:00:00 2001 From: "James R. Wilcox" Date: Thu, 31 Aug 2023 16:48:14 -0700 Subject: [PATCH 1/5] Bump egglog commit hash --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index f7ba00c74..0b31fa82f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,7 @@ name = "files" [dependencies] -egglog = { git = "https://github.com/egraphs-good/egglog", rev = "a4768b1" } +egglog = { git = "https://github.com/egraphs-good/egglog", rev = "c83fc75" } log = "0.4.19" thiserror = "1" lalrpop-util = { version = "0.19.8", features = ["lexer"] } From 6c4787c90aa3058f351f0d8ff3df610a11429846 Mon Sep 17 00:00:00 2001 From: "James R. Wilcox" Date: Wed, 23 Aug 2023 15:19:10 -0700 Subject: [PATCH 2/5] Use terms instead of exprs in CFG conversion (but no memo yet) --- src/conversions.rs | 508 ++++++++++++++++++++++----------------------- src/lib.rs | 11 +- 2 files changed, 253 insertions(+), 266 deletions(-) diff --git a/src/conversions.rs b/src/conversions.rs index 9e4336fba..48f928c4a 100644 --- a/src/conversions.rs +++ b/src/conversions.rs @@ -9,47 +9,44 @@ use crate::{ }; use bril_rs::{Argument, Code, EffectOps, Instruction, Literal, Program, Type, ValueOps}; use egglog::ast::{Expr, Symbol}; +use egglog::{match_term_app, Term, TermDag}; use ordered_float::OrderedFloat; impl Optimizer { - /// Convert an egglog expression back into a StructuredFunction - pub(crate) fn expr_to_structured_func(&mut self, expr: Expr) -> StructuredFunction { - if let Expr::Call(func, args) = expr { - assert_eq!(func.to_string(), "Func"); - match &args.as_slice() { - [func_name, argslist, body] => { - let args = Self::conslist_to_vec(argslist, "Arg") - .into_iter() - .map(|arg| Self::expr_to_argument(&arg)) - .collect(); - if let Expr::Lit(egglog::ast::Literal::String(fname)) = func_name { - StructuredFunction { - name: fname.to_string(), - args, - block: self.expr_to_structured_block(body), - } - } else { - panic!("expected string literal for func name"); - } + pub(crate) fn term_to_structured_func( + &mut self, + termdag: &TermDag, + t: &Term, + ) -> StructuredFunction { + match_term_app!(t; { + ("Func", [func_name, argslist, body]) => { + let args = self.term_conslist_to_vec(termdag, &termdag.get(*argslist), "Arg") + .into_iter() + .map(|arg| Self::term_to_argument(termdag, &arg)) + .collect(); + + let fname = Self::string_term_to_string(&termdag.get(*func_name)); + StructuredFunction { + name: fname.to_string(), + args, + block: self.term_to_structured_block(termdag, &termdag.get(*body)), } - _ => panic!("expected 2 args in expr_to_func"), } - } else { - panic!("expected call in expr_to_func"); - } + (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) + }) } - fn expr_to_argument(expr: &Expr) -> Argument { - let Expr::Call(op, args) = expr else { - panic!("expected call in expr_to_argument"); - }; - match (op.as_str(), args.as_slice()) { - ("Arg", [Expr::Lit(egglog::ast::Literal::String(name)), ty]) => Argument { - name: name.to_string(), - arg_type: Self::expr_to_type(ty), - }, - _ => panic!("unknown argument"), - } + fn term_to_argument(termdag: &TermDag, term: &Term) -> Argument { + match_term_app!(term; { + ("Arg", [name, ty]) => { + let name = Self::string_term_to_string(&termdag.get(*name)); + Argument { + name: name.to_string(), + arg_type: Self::term_to_type(termdag, &termdag.get(*ty)), + } + } + (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) + }) } fn argument_to_expr(arg: &Argument) -> Expr { @@ -62,104 +59,100 @@ impl Optimizer { ) } - pub(crate) fn expr_to_structured_block(&mut self, expr: &Expr) -> StructuredBlock { - if let Expr::Call(func, args) = expr { - match (func.as_str(), &args.as_slice()) { - ("Block", [block]) => { - StructuredBlock::Block(Box::new(self.expr_to_structured_block(block))) - } - ("Basic", [basic_block]) => { - StructuredBlock::Basic(Box::new(self.expr_to_basic_block(basic_block))) - } - ("Ite", [name, then_branch, else_branch]) => StructuredBlock::Ite( - Self::string_expr_to_string(name), - Box::new(self.expr_to_structured_block(then_branch)), - Box::new(self.expr_to_structured_block(else_branch)), - ), - ("Loop", [block]) => { - StructuredBlock::Loop(Box::new(self.expr_to_structured_block(block))) - } - ("Sequence", [block, rest]) => StructuredBlock::Sequence(vec![ - self.expr_to_structured_block(block), - self.expr_to_structured_block(rest), - ]), - ("Break", [n]) => { - if let Expr::Lit(egglog::ast::Literal::Int(n)) = n { - StructuredBlock::Break((*n).try_into().unwrap()) - } else { - panic!("expected int literal for break"); - } + pub(crate) fn term_to_structured_block( + &mut self, + termdag: &TermDag, + term: &Term, + ) -> StructuredBlock { + match_term_app!(term; { + ("Block", [block]) => { + StructuredBlock::Block(Box::new(self.term_to_structured_block(termdag, &termdag.get(*block)))) + }, + ("Basic", [basic_block]) => { + StructuredBlock::Basic(Box::new(self.term_to_basic_block(termdag, &termdag.get(*basic_block)))) + }, + ("Ite", [name, then_branch, else_branch]) => { + let string = Self::string_term_to_string(&termdag.get(*name)); + StructuredBlock::Ite( + string.to_string(), + Box::new(self.term_to_structured_block(termdag, &termdag.get(*then_branch))), + Box::new(self.term_to_structured_block(termdag, &termdag.get(*else_branch))), + ) + }, + ("Loop", [block]) => { + StructuredBlock::Loop(Box::new(self.term_to_structured_block(termdag, &termdag.get(*block)))) + }, + ("Sequence", [block, rest]) => StructuredBlock::Sequence(vec![ + self.term_to_structured_block(termdag, &termdag.get(*block)), + self.term_to_structured_block(termdag, &termdag.get(*rest)), + ]), + ("Break", [n]) => { + if let Term::Lit(egglog::ast::Literal::Int(n)) = termdag.get(*n) { + StructuredBlock::Break(n.try_into().unwrap()) + } else { + panic!("expected int literal for break"); } - ("Return", [val]) => { - if let Expr::Call(op, args) = val { - match op.as_str() { - "Void" => StructuredBlock::Return(None), - "ReturnValue" => { - assert_eq!(args.len(), 1); - match &args[0] { - Expr::Lit(egglog::ast::Literal::String(val)) => { - StructuredBlock::Return(Some(val.to_string())) - } - _ => panic!("expected string literal for return value"), - } + }, + ("Return", [val]) => { + match_term_app!(termdag.get(*val); { + ("Void", _) => StructuredBlock::Return(None), + ("ReturnValue", [arg]) => { + match termdag.get(*arg) { + Term::Lit(egglog::ast::Literal::String(s)) => { + StructuredBlock::Return(Some(s.to_string())) } - _ => panic!("expected void or return value"), + _ => panic!("expected string literal for return value"), } - } else { - panic!("expected call for return"); } - } - _ => panic!("unknown structured block"), + (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) + }) } - } else { - panic!("expected call in expr_to_structured_block"); - } + (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) + }) } - pub(crate) fn expr_to_basic_block(&mut self, expr: &Expr) -> BasicBlock { - if let Expr::Call(op, args) = expr { - assert_eq!(op.as_str(), "BlockNamed"); - - match &args.as_slice() { - [Expr::Lit(egglog::ast::Literal::String(name)), code] => { - let code_vec = Self::conslist_to_vec(code, "Code"); - let mut instrs = vec![]; + pub(crate) fn term_to_basic_block(&mut self, termdag: &TermDag, term: &Term) -> BasicBlock { + match_term_app!(term; { + ("BlockNamed", [name, code]) => { + let name = Self::string_term_to_string(&termdag.get(*name)); + let code_vec = self.term_conslist_to_vec(termdag, &termdag.get(*code), "Code"); + let mut instrs = vec![]; + // let mut memo = HashMap::::new(); - for expr in code_vec { - self.expr_to_instructions(&expr, &mut instrs); - } + for t in code_vec { + self.term_to_instructions(termdag, &t, &mut instrs); + } - BasicBlock { - name: BlockName::Named(name.to_string()), - footer: Default::default(), - instrs, - pos: None, - } + BasicBlock { + name: BlockName::Named(name.to_string()), + footer: Default::default(), + instrs, + pos: None, } - _ => panic!("expected 2 args in expr_to_basic_block"), } - } else { - panic!("expected call in expr_to_basic_block"); - } + (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) + }) } - fn conslist_to_vec_helper(expr: &Expr, res: &mut Vec, prefix: &str) { - match expr { - Expr::Call(op, args) => match (op.as_str(), args.as_slice()) { - (op, [head, tail]) if op == prefix.to_string() + "Cons" => { - res.push(head.clone()); - Self::conslist_to_vec_helper(tail, res, prefix); - } - (op, []) if op == prefix.to_string() + "Nil" => {} - _ => panic!("expected Cons or Nil"), + fn term_conslist_to_vec_helper( + termdag: &TermDag, + term: &Term, + res: &mut Vec, + prefix: &str, + ) { + match_term_app!(term; { + (op, [head, tail]) if op == prefix.to_string() + "Cons" => { + res.push(termdag.get(*head)); + Self::term_conslist_to_vec_helper(termdag, &termdag.get(*tail), res, prefix); }, - _ => panic!("expected call in codelist_to_vec"), - } + (op, []) if op == prefix.to_string() + "Nil" => {} + (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) + }) } - fn conslist_to_vec(expr: &Expr, prefix: &str) -> Vec { + fn term_conslist_to_vec(&self, termdag: &TermDag, term: &Term, prefix: &str) -> Vec { let mut res = vec![]; - Self::conslist_to_vec_helper(expr, &mut res, prefix); + Self::term_conslist_to_vec_helper(termdag, term, &mut res, prefix); res } @@ -171,77 +164,61 @@ impl Optimizer { current } - fn expr_to_instructions(&mut self, expr: &Expr, res: &mut Vec) { - if let Expr::Call(op, args) = expr { - match op.to_string().as_str() { - "Print" => { - assert!(args.len() == 1); - let arg = self.expr_to_code(&args[0], res, None); - - res.push(Instruction::Effect { - op: EffectOps::Print, - args: vec![arg], - funcs: vec![], - labels: vec![], - pos: None, - }); - } - "End" => { - assert!(args.is_empty()); - } - "Assign" => { - assert!( - args.len() == 2, - "expected 2 args in Assign. got: {:?}", - args - ); - let dest = match &args[0] { - Expr::Lit(egglog::ast::Literal::String(dest)) => dest.to_string(), - _ => panic!("expected string literal for dest"), - }; - self.expr_to_code(&args[1], res, Some(dest)); - } - "store" | "free" => { - let args = args - .iter() - .map(|arg| self.expr_to_code(arg, res, None)) - .collect::>(); - - res.push(Instruction::Effect { - op: serde_json::from_str(&format!("\"{}\"", op)).unwrap(), - args, - funcs: vec![], - labels: vec![], - pos: None, - }); - } - "alloc" => { - assert!(args.len() == 3); - let atype = Self::expr_to_type(&args[0]); - let dest = Self::string_expr_to_string(&args[1]); - let arg = self.expr_to_code(&args[2], res, None); - res.push(Instruction::Value { - dest, - args: vec![arg], - funcs: vec![], - op: ValueOps::Alloc, - labels: vec![], - pos: None, - op_type: atype, - }); - } + fn term_to_instructions(&mut self, termdag: &TermDag, term: &Term, res: &mut Vec) { + match_term_app!(term; { + ("Print", [arg]) => { + let arg = self.term_to_code(termdag, &termdag.get(*arg), res, None); - _ => panic!("unknown effect in body_to_code {}", op), - } - } else { - panic!("expected call in body_to_code"); - } + res.push(Instruction::Effect { + op: EffectOps::Print, + args: vec![arg], + funcs: vec![], + labels: vec![], + pos: None, + }); + }, + ("End", []) => {}, + ("Assign", [dest, src]) => { + let dest = Self::string_term_to_string(&termdag.get(*dest)); + self.term_to_code(termdag, &termdag.get(*src), res, Some(dest.to_string())); + }, + (op @ ("store" | "free"), args) => { + let args = args + .iter() + .map(|arg| self.term_to_code(termdag, &termdag.get(*arg), res, None)) + .collect::>(); + + + res.push(Instruction::Effect { + op: serde_json::from_str(&format!("\"{}\"", op)).unwrap(), + args, + funcs: vec![], + labels: vec![], + pos: None, + }); + }, + ("alloc", [atype, dest, arg]) => { + let atype = Self::term_to_type(termdag, &termdag.get(*atype)); + let dest = Self::string_term_to_string(&termdag.get(*dest)); + let arg = self.term_to_code(termdag, &termdag.get(*arg), res, None); + res.push(Instruction::Value { + dest, + args: vec![arg], + funcs: vec![], + op: ValueOps::Alloc, + labels: vec![], + pos: None, + op_type: atype, + }); + }, + (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) + }) } - // TODO memoize exprs for common subexpression elimination - pub(crate) fn expr_to_code( + pub(crate) fn term_to_code( &mut self, - expr: &Expr, + termdag: &TermDag, + term: &Term, res: &mut Vec, assign_to: Option, ) -> String { @@ -249,8 +226,9 @@ impl Optimizer { Some(dest) => dest.clone(), None => self.fresh_var(), }; - match expr { - Expr::Lit(literal) => { + + match term { + Term::Lit(literal) => { res.push(Instruction::Constant { dest: dest.clone(), op: bril_rs::ConstOps::Const, @@ -260,80 +238,91 @@ impl Optimizer { }); dest } - Expr::Var(var) => { + Term::Var(var) => { if let Some(_output) = assign_to { panic!("Cannot assign var to var") } else { var.to_string() } } - Expr::Call(op, args) => match op.to_string().as_str() { - "Var" => { - assert!(args.len() == 1); - match &args[0] { - Expr::Lit(egglog::ast::Literal::String(var)) => var.to_string(), - _ => panic!("expected string literal for var"), - } - } - "ReturnValue" => self.expr_to_code(&args[0], res, assign_to), - "Int" | "True" | "False" => { - let literal = match op.to_string().as_str() { - "Int" => Literal::Int(args[1].to_string().parse().unwrap()), - "True" => Literal::Bool(true), - "False" => Literal::Bool(false), - "Char" => { - assert_eq!(args[0].to_string().len(), 1); - Literal::Char(args[0].to_string().chars().next().unwrap()) + _ => { + match_term_app!(term; { + ("Var", [arg]) => { + match termdag.get(*arg) { + Term::Lit(egglog::ast::Literal::String(var)) => var.to_string(), + _ => panic!("expected string literal for var"), } - _ => panic!("unknown literal"), - }; - res.push(Instruction::Constant { - dest: dest.clone(), - op: bril_rs::ConstOps::Const, - value: literal, - pos: None, - const_type: Self::expr_to_type(&args[0]), - }); - dest - } - "phi" => { - assert!(args.len() == 5); - let etype = Self::expr_to_type(&args[0]); - let arg1 = self.expr_to_code(&args[1], res, None); - let arg2 = self.expr_to_code(&args[2], res, None); - let label1 = Self::string_expr_to_string(&args[3]); - let label2 = Self::string_expr_to_string(&args[4]); - res.push(Instruction::Value { - dest: dest.clone(), - args: vec![arg1, arg2], - funcs: vec![], - op: ValueOps::Phi, - labels: vec![label1, label2], - pos: None, - op_type: etype, - }); - dest - } - _ => { - assert!(op.as_str() != "Void"); - let etype = Self::expr_to_type(&args[0]); - let args_vars = args - .iter() - .skip(1) - .map(|arg| self.expr_to_code(arg, res, None)) - .collect::>(); - res.push(Instruction::Value { - dest: dest.clone(), - args: args_vars, - funcs: vec![], - op: egglog_op_to_bril(*op), - labels: vec![], - pos: None, - op_type: etype, - }); - dest - } - }, + }, + ("ReturnValue", [arg]) => self.term_to_code(termdag, &termdag.get(*arg), res, assign_to), + (op @ ("True" | "False" | "Int" | "Float" | "Char"), [ty, args @ ..]) => { + let lit = match (op, args) { + ("True", []) => Literal::Bool(true), + ("False", []) => Literal::Bool(false), + ("Int", [arg]) => { + let arg = termdag.get(*arg); + let arg_s = termdag.to_string(&arg); + Literal::Int(arg_s.parse::().unwrap()) + } + ("Float", [arg]) => { + let arg = termdag.get(*arg); + let arg_s = termdag.to_string(&arg); + Literal::Float(arg_s.parse::().unwrap()) + } + ("Char", [arg]) => { + let arg = termdag.get(*arg); + let arg_s = termdag.to_string(&arg); + assert_eq!(arg_s.len(), 1); + Literal::Char(arg_s.chars().next().unwrap()) + } + _ => panic!("unexpected args to literal in term_to_code") + }; + res.push(Instruction::Constant { + dest: dest.clone(), + op: bril_rs::ConstOps::Const, + value: lit, + pos: None, + const_type: Self::term_to_type(termdag, &termdag.get(*ty)), + }); + dest + }, + ("phi", [etype, arg1, arg2, label1, label2]) => { + let etype = Self::term_to_type(termdag, &termdag.get(*etype)); + let arg1 = self.term_to_code(termdag, &termdag.get(*arg1), res, None); + let arg2 = self.term_to_code(termdag, &termdag.get(*arg2), res, None); + let label1 = Self::string_term_to_string(&termdag.get(*label1)); + let label2 = Self::string_term_to_string(&termdag.get(*label2)); + res.push(Instruction::Value { + dest: dest.clone(), + args: vec![arg1, arg2], + funcs: vec![], + op: ValueOps::Phi, + labels: vec![label1, label2], + pos: None, + op_type: etype, + }); + dest + }, + (op, args) => { + assert!(op != "Void"); + let etype = Self::term_to_type(termdag, &termdag.get(args[0])); + let args_vars = args + .iter() + .skip(1) + .map(|arg| self.term_to_code(termdag, &termdag.get(*arg), res, None)) + .collect::>(); + res.push(Instruction::Value { + dest: dest.clone(), + args: args_vars, + funcs: vec![], + op: egglog_op_to_bril(op.into()), + labels: vec![], + pos: None, + op_type: etype, + }); + dest + } + }) + } } } @@ -412,8 +401,8 @@ impl Optimizer { Expr::Call("Var".into(), vec![self.string_to_expr(string)]) } - fn string_expr_to_string(expr: &Expr) -> String { - if let Expr::Lit(egglog::ast::Literal::String(string)) = expr { + fn string_term_to_string(term: &Term) -> String { + if let Term::Lit(egglog::ast::Literal::String(string)) = term { string.to_string() } else { panic!("expected string literal"); @@ -622,18 +611,15 @@ impl Optimizer { } } - pub(crate) fn expr_to_type(expr: &Expr) -> Type { - let Expr::Call(op, args) = expr else { - panic!("expected call in expr_to_type"); - }; - match (op.as_str(), args.as_slice()) { + pub(crate) fn term_to_type(termdag: &TermDag, term: &Term) -> Type { + match_term_app!(term; { ("IntT", []) => Type::Int, ("BoolT", []) => Type::Bool, ("FloatT", []) => Type::Float, ("CharT", []) => Type::Char, - ("PointerT", [child]) => Type::Pointer(Box::new(Self::expr_to_type(child))), - _ => panic!("unknown type"), - } + ("PointerT", [child]) => Type::Pointer(Box::new(Self::term_to_type(termdag, &termdag.get(*child)))), + (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) + }) } pub(crate) fn pretty_print_expr(expr: &Expr) -> String { diff --git a/src/lib.rs b/src/lib.rs index f6cade3a3..4b7ff7939 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -239,14 +239,15 @@ impl Optimizer { let mut keys = egg_fns.keys().collect::>(); keys.sort(); + let mut termdag = Default::default(); let mut result = vec![]; for name in keys { let expr = egg_fns.get(name).unwrap(); - let mut rep = egraph - .extract_expr(expr.clone(), 0) + let (sort, value) = egraph + .eval_expr(expr, None, true) .map_err(EggCCError::EggLog)?; - let extracted = rep.termdag.term_to_expr(&rep.expr); - let structured_func = self.expr_to_structured_func(extracted); + let (_cost, term) = egraph.extract(value, &mut termdag, &sort); + let structured_func = self.term_to_structured_func(&termdag, &term); result.push(structured_func); } @@ -285,7 +286,7 @@ impl Optimizer { (lt Type Expr Expr) (ptradd Type Expr Expr) (load Type Expr) - + ) (datatype RetVal From 2225f60b8385742f25ecccc0b64b42c70b9070cf Mon Sep 17 00:00:00 2001 From: "James R. Wilcox" Date: Wed, 23 Aug 2023 15:55:50 -0700 Subject: [PATCH 3/5] conversions.rs: refactor to introduce TermConverter Avoids passing termdag all over the place and prepares for memo table. --- src/conversions.rs | 205 +++++++++++++++++++++++---------------------- 1 file changed, 105 insertions(+), 100 deletions(-) diff --git a/src/conversions.rs b/src/conversions.rs index 48f928c4a..e95a0f608 100644 --- a/src/conversions.rs +++ b/src/conversions.rs @@ -12,92 +12,79 @@ use egglog::ast::{Expr, Symbol}; use egglog::{match_term_app, Term, TermDag}; use ordered_float::OrderedFloat; -impl Optimizer { - pub(crate) fn term_to_structured_func( - &mut self, - termdag: &TermDag, - t: &Term, - ) -> StructuredFunction { +pub(crate) struct TermConverter<'a> { + optimizer: &'a mut Optimizer, + termdag: &'a TermDag, +} + +impl TermConverter<'_> { + pub(crate) fn term_to_structured_func(&mut self, t: &Term) -> StructuredFunction { match_term_app!(t; { ("Func", [func_name, argslist, body]) => { - let args = self.term_conslist_to_vec(termdag, &termdag.get(*argslist), "Arg") + let args = self.term_conslist_to_vec(&self.termdag.get(*argslist), "Arg") .into_iter() - .map(|arg| Self::term_to_argument(termdag, &arg)) + .map(|arg| self.term_to_argument(&arg)) .collect(); - let fname = Self::string_term_to_string(&termdag.get(*func_name)); + let fname = Optimizer::string_term_to_string(&self.termdag.get(*func_name)); StructuredFunction { name: fname.to_string(), args, - block: self.term_to_structured_block(termdag, &termdag.get(*body)), + block: self.term_to_structured_block(&self.termdag.get(*body)), } } (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) }) } - fn term_to_argument(termdag: &TermDag, term: &Term) -> Argument { + fn term_to_argument(&self, term: &Term) -> Argument { match_term_app!(term; { ("Arg", [name, ty]) => { - let name = Self::string_term_to_string(&termdag.get(*name)); + let name = Optimizer::string_term_to_string(&self.termdag.get(*name)); Argument { name: name.to_string(), - arg_type: Self::term_to_type(termdag, &termdag.get(*ty)), + arg_type: self.term_to_type(&self.termdag.get(*ty)), } } (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) }) } - fn argument_to_expr(arg: &Argument) -> Expr { - Expr::Call( - "Arg".into(), - vec![ - Expr::Lit(egglog::ast::Literal::String(arg.name.clone().into())), - Self::type_to_expr(&arg.arg_type), - ], - ) - } - - pub(crate) fn term_to_structured_block( - &mut self, - termdag: &TermDag, - term: &Term, - ) -> StructuredBlock { + pub(crate) fn term_to_structured_block(&mut self, term: &Term) -> StructuredBlock { match_term_app!(term; { ("Block", [block]) => { - StructuredBlock::Block(Box::new(self.term_to_structured_block(termdag, &termdag.get(*block)))) + StructuredBlock::Block(Box::new(self.term_to_structured_block(&self.termdag.get(*block)))) }, ("Basic", [basic_block]) => { - StructuredBlock::Basic(Box::new(self.term_to_basic_block(termdag, &termdag.get(*basic_block)))) + StructuredBlock::Basic(Box::new(self.term_to_basic_block(&self.termdag.get(*basic_block)))) }, ("Ite", [name, then_branch, else_branch]) => { - let string = Self::string_term_to_string(&termdag.get(*name)); + let string = Optimizer::string_term_to_string(&self.termdag.get(*name)); StructuredBlock::Ite( string.to_string(), - Box::new(self.term_to_structured_block(termdag, &termdag.get(*then_branch))), - Box::new(self.term_to_structured_block(termdag, &termdag.get(*else_branch))), + Box::new(self.term_to_structured_block(&self.termdag.get(*then_branch))), + Box::new(self.term_to_structured_block(&self.termdag.get(*else_branch))), ) }, ("Loop", [block]) => { - StructuredBlock::Loop(Box::new(self.term_to_structured_block(termdag, &termdag.get(*block)))) + StructuredBlock::Loop(Box::new(self.term_to_structured_block(&self.termdag.get(*block)))) }, ("Sequence", [block, rest]) => StructuredBlock::Sequence(vec![ - self.term_to_structured_block(termdag, &termdag.get(*block)), - self.term_to_structured_block(termdag, &termdag.get(*rest)), + self.term_to_structured_block(&self.termdag.get(*block)), + self.term_to_structured_block(&self.termdag.get(*rest)), ]), ("Break", [n]) => { - if let Term::Lit(egglog::ast::Literal::Int(n)) = termdag.get(*n) { + if let Term::Lit(egglog::ast::Literal::Int(n)) = self.termdag.get(*n) { StructuredBlock::Break(n.try_into().unwrap()) } else { panic!("expected int literal for break"); } }, ("Return", [val]) => { - match_term_app!(termdag.get(*val); { + match_term_app!(self.termdag.get(*val); { ("Void", _) => StructuredBlock::Return(None), ("ReturnValue", [arg]) => { - match termdag.get(*arg) { + match self.termdag.get(*arg) { Term::Lit(egglog::ast::Literal::String(s)) => { StructuredBlock::Return(Some(s.to_string())) } @@ -111,16 +98,16 @@ impl Optimizer { }) } - pub(crate) fn term_to_basic_block(&mut self, termdag: &TermDag, term: &Term) -> BasicBlock { + pub(crate) fn term_to_basic_block(&mut self, term: &Term) -> BasicBlock { match_term_app!(term; { ("BlockNamed", [name, code]) => { - let name = Self::string_term_to_string(&termdag.get(*name)); - let code_vec = self.term_conslist_to_vec(termdag, &termdag.get(*code), "Code"); + let name = Optimizer::string_term_to_string(&self.termdag.get(*name)); + let code_vec = self.term_conslist_to_vec(&self.termdag.get(*code), "Code"); let mut instrs = vec![]; // let mut memo = HashMap::::new(); for t in code_vec { - self.term_to_instructions(termdag, &t, &mut instrs); + self.term_to_instructions(&t, &mut instrs); } BasicBlock { @@ -134,40 +121,27 @@ impl Optimizer { }) } - fn term_conslist_to_vec_helper( - termdag: &TermDag, - term: &Term, - res: &mut Vec, - prefix: &str, - ) { + fn term_conslist_to_vec_helper(&self, term: &Term, res: &mut Vec, prefix: &str) { match_term_app!(term; { (op, [head, tail]) if op == prefix.to_string() + "Cons" => { - res.push(termdag.get(*head)); - Self::term_conslist_to_vec_helper(termdag, &termdag.get(*tail), res, prefix); + res.push(self.termdag.get(*head)); + self.term_conslist_to_vec_helper(&self.termdag.get(*tail), res, prefix); }, (op, []) if op == prefix.to_string() + "Nil" => {} (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) }) } - fn term_conslist_to_vec(&self, termdag: &TermDag, term: &Term, prefix: &str) -> Vec { + fn term_conslist_to_vec(&self, term: &Term, prefix: &str) -> Vec { let mut res = vec![]; - Self::term_conslist_to_vec_helper(termdag, term, &mut res, prefix); + self.term_conslist_to_vec_helper(term, &mut res, prefix); res } - fn vec_to_cons_list(vec: Vec, prefix: &str) -> Expr { - let mut current = Expr::Call(format!("{prefix}Nil").into(), vec![]); - for expr in vec.into_iter().rev() { - current = Expr::Call(format!("{prefix}Cons").into(), vec![expr, current]); - } - current - } - - fn term_to_instructions(&mut self, termdag: &TermDag, term: &Term, res: &mut Vec) { + fn term_to_instructions(&mut self, term: &Term, res: &mut Vec) { match_term_app!(term; { ("Print", [arg]) => { - let arg = self.term_to_code(termdag, &termdag.get(*arg), res, None); + let arg = self.term_to_code(&self.termdag.get(*arg), res, None); res.push(Instruction::Effect { op: EffectOps::Print, @@ -179,13 +153,13 @@ impl Optimizer { }, ("End", []) => {}, ("Assign", [dest, src]) => { - let dest = Self::string_term_to_string(&termdag.get(*dest)); - self.term_to_code(termdag, &termdag.get(*src), res, Some(dest.to_string())); + let dest = Optimizer::string_term_to_string(&self.termdag.get(*dest)); + self.term_to_code(&self.termdag.get(*src), res, Some(dest.to_string())); }, (op @ ("store" | "free"), args) => { let args = args .iter() - .map(|arg| self.term_to_code(termdag, &termdag.get(*arg), res, None)) + .map(|arg| self.term_to_code(&self.termdag.get(*arg), res, None)) .collect::>(); @@ -198,9 +172,9 @@ impl Optimizer { }); }, ("alloc", [atype, dest, arg]) => { - let atype = Self::term_to_type(termdag, &termdag.get(*atype)); - let dest = Self::string_term_to_string(&termdag.get(*dest)); - let arg = self.term_to_code(termdag, &termdag.get(*arg), res, None); + let atype = self.term_to_type(&self.termdag.get(*atype)); + let dest = Optimizer::string_term_to_string(&self.termdag.get(*dest)); + let arg = self.term_to_code(&self.termdag.get(*arg), res, None); res.push(Instruction::Value { dest, args: vec![arg], @@ -217,14 +191,13 @@ impl Optimizer { pub(crate) fn term_to_code( &mut self, - termdag: &TermDag, term: &Term, res: &mut Vec, assign_to: Option, ) -> String { let dest = match &assign_to { Some(dest) => dest.clone(), - None => self.fresh_var(), + None => self.optimizer.fresh_var(), }; match term { @@ -232,9 +205,9 @@ impl Optimizer { res.push(Instruction::Constant { dest: dest.clone(), op: bril_rs::ConstOps::Const, - value: self.literal_to_bril(literal), + value: self.optimizer.literal_to_bril(literal), pos: None, - const_type: self.literal_to_type(literal), + const_type: self.optimizer.literal_to_type(literal), }); dest } @@ -248,29 +221,29 @@ impl Optimizer { _ => { match_term_app!(term; { ("Var", [arg]) => { - match termdag.get(*arg) { + match self.termdag.get(*arg) { Term::Lit(egglog::ast::Literal::String(var)) => var.to_string(), _ => panic!("expected string literal for var"), } }, - ("ReturnValue", [arg]) => self.term_to_code(termdag, &termdag.get(*arg), res, assign_to), + ("ReturnValue", [arg]) => self.term_to_code(&self.termdag.get(*arg), res, assign_to), (op @ ("True" | "False" | "Int" | "Float" | "Char"), [ty, args @ ..]) => { let lit = match (op, args) { ("True", []) => Literal::Bool(true), ("False", []) => Literal::Bool(false), ("Int", [arg]) => { - let arg = termdag.get(*arg); - let arg_s = termdag.to_string(&arg); + let arg = self.termdag.get(*arg); + let arg_s = self.termdag.to_string(&arg); Literal::Int(arg_s.parse::().unwrap()) } ("Float", [arg]) => { - let arg = termdag.get(*arg); - let arg_s = termdag.to_string(&arg); + let arg = self.termdag.get(*arg); + let arg_s = self.termdag.to_string(&arg); Literal::Float(arg_s.parse::().unwrap()) } ("Char", [arg]) => { - let arg = termdag.get(*arg); - let arg_s = termdag.to_string(&arg); + let arg = self.termdag.get(*arg); + let arg_s = self.termdag.to_string(&arg); assert_eq!(arg_s.len(), 1); Literal::Char(arg_s.chars().next().unwrap()) } @@ -281,16 +254,16 @@ impl Optimizer { op: bril_rs::ConstOps::Const, value: lit, pos: None, - const_type: Self::term_to_type(termdag, &termdag.get(*ty)), + const_type: self.term_to_type(&self.termdag.get(*ty)), }); dest }, ("phi", [etype, arg1, arg2, label1, label2]) => { - let etype = Self::term_to_type(termdag, &termdag.get(*etype)); - let arg1 = self.term_to_code(termdag, &termdag.get(*arg1), res, None); - let arg2 = self.term_to_code(termdag, &termdag.get(*arg2), res, None); - let label1 = Self::string_term_to_string(&termdag.get(*label1)); - let label2 = Self::string_term_to_string(&termdag.get(*label2)); + let etype = self.term_to_type(&self.termdag.get(*etype)); + let arg1 = self.term_to_code(&self.termdag.get(*arg1), res, None); + let arg2 = self.term_to_code(&self.termdag.get(*arg2), res, None); + let label1 = Optimizer::string_term_to_string(&self.termdag.get(*label1)); + let label2 = Optimizer::string_term_to_string(&self.termdag.get(*label2)); res.push(Instruction::Value { dest: dest.clone(), args: vec![arg1, arg2], @@ -304,11 +277,11 @@ impl Optimizer { }, (op, args) => { assert!(op != "Void"); - let etype = Self::term_to_type(termdag, &termdag.get(args[0])); + let etype = self.term_to_type(&self.termdag.get(args[0])); let args_vars = args .iter() .skip(1) - .map(|arg| self.term_to_code(termdag, &termdag.get(*arg), res, None)) + .map(|arg| self.term_to_code(&self.termdag.get(*arg), res, None)) .collect::>(); res.push(Instruction::Value { dest: dest.clone(), @@ -326,6 +299,31 @@ impl Optimizer { } } + pub(crate) fn term_to_type(&self, term: &Term) -> Type { + match_term_app!(term; { + ("IntT", []) => Type::Int, + ("BoolT", []) => Type::Bool, + ("FloatT", []) => Type::Float, + ("CharT", []) => Type::Char, + ("PointerT", [child]) => Type::Pointer(Box::new(self.term_to_type(&self.termdag.get(*child)))), + (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) + }) + } +} + +impl Optimizer { + pub(crate) fn term_to_structured_func( + &mut self, + termdag: &TermDag, + term: &Term, + ) -> StructuredFunction { + let mut converter = TermConverter { + optimizer: self, + termdag, + }; + converter.term_to_structured_func(term) + } + pub(crate) fn func_to_expr(&mut self, func: &StructuredFunction) -> Expr { let arg_exprs = func .args @@ -343,6 +341,24 @@ impl Optimizer { ) } + fn vec_to_cons_list(vec: Vec, prefix: &str) -> Expr { + let mut current = Expr::Call(format!("{prefix}Nil").into(), vec![]); + for expr in vec.into_iter().rev() { + current = Expr::Call(format!("{prefix}Cons").into(), vec![expr, current]); + } + current + } + + fn argument_to_expr(arg: &Argument) -> Expr { + Expr::Call( + "Arg".into(), + vec![ + Expr::Lit(egglog::ast::Literal::String(arg.name.clone().into())), + Self::type_to_expr(&arg.arg_type), + ], + ) + } + pub(crate) fn structured_block_to_expr(&mut self, structured_block: &StructuredBlock) -> Expr { match structured_block { StructuredBlock::Ite(var, then, els) => Expr::Call( @@ -611,17 +627,6 @@ impl Optimizer { } } - pub(crate) fn term_to_type(termdag: &TermDag, term: &Term) -> Type { - match_term_app!(term; { - ("IntT", []) => Type::Int, - ("BoolT", []) => Type::Bool, - ("FloatT", []) => Type::Float, - ("CharT", []) => Type::Char, - ("PointerT", [child]) => Type::Pointer(Box::new(Self::term_to_type(termdag, &termdag.get(*child)))), - (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) - }) - } - pub(crate) fn pretty_print_expr(expr: &Expr) -> String { Self::pretty_print_expr_with_acc(expr, 0) } From d02952a8f77d8ba54e3462a4261e08f4e28e9b3d Mon Sep 17 00:00:00 2001 From: "James R. Wilcox" Date: Wed, 23 Aug 2023 16:17:34 -0700 Subject: [PATCH 4/5] Use TermIds everywhere to cleanup recursion --- src/conversions.rs | 150 +++++++++++++++++++++++---------------------- 1 file changed, 77 insertions(+), 73 deletions(-) diff --git a/src/conversions.rs b/src/conversions.rs index e95a0f608..edc7c5854 100644 --- a/src/conversions.rs +++ b/src/conversions.rs @@ -9,7 +9,7 @@ use crate::{ }; use bril_rs::{Argument, Code, EffectOps, Instruction, Literal, Program, Type, ValueOps}; use egglog::ast::{Expr, Symbol}; -use egglog::{match_term_app, Term, TermDag}; +use egglog::{match_term_app, Term, TermDag, TermId}; use ordered_float::OrderedFloat; pub(crate) struct TermConverter<'a> { @@ -18,73 +18,77 @@ pub(crate) struct TermConverter<'a> { } impl TermConverter<'_> { - pub(crate) fn term_to_structured_func(&mut self, t: &Term) -> StructuredFunction { - match_term_app!(t; { + pub(crate) fn get(&self, id: &TermId) -> Term { + self.termdag.get(*id) + } + + pub(crate) fn term_to_structured_func(&mut self, id: &TermId) -> StructuredFunction { + match_term_app!(self.get(id); { ("Func", [func_name, argslist, body]) => { - let args = self.term_conslist_to_vec(&self.termdag.get(*argslist), "Arg") + let args = self.term_conslist_to_vec(argslist, "Arg") .into_iter() .map(|arg| self.term_to_argument(&arg)) .collect(); - let fname = Optimizer::string_term_to_string(&self.termdag.get(*func_name)); + let fname = self.string_term_to_string(func_name); StructuredFunction { name: fname.to_string(), args, - block: self.term_to_structured_block(&self.termdag.get(*body)), + block: self.term_to_structured_block(body), } } (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) }) } - fn term_to_argument(&self, term: &Term) -> Argument { - match_term_app!(term; { + fn term_to_argument(&self, id: &TermId) -> Argument { + match_term_app!(self.get(id); { ("Arg", [name, ty]) => { - let name = Optimizer::string_term_to_string(&self.termdag.get(*name)); + let name = self.string_term_to_string(name); Argument { name: name.to_string(), - arg_type: self.term_to_type(&self.termdag.get(*ty)), + arg_type: self.term_to_type(ty), } } (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) }) } - pub(crate) fn term_to_structured_block(&mut self, term: &Term) -> StructuredBlock { - match_term_app!(term; { + pub(crate) fn term_to_structured_block(&mut self, id: &TermId) -> StructuredBlock { + match_term_app!(self.get(id); { ("Block", [block]) => { - StructuredBlock::Block(Box::new(self.term_to_structured_block(&self.termdag.get(*block)))) + StructuredBlock::Block(Box::new(self.term_to_structured_block(block))) }, ("Basic", [basic_block]) => { - StructuredBlock::Basic(Box::new(self.term_to_basic_block(&self.termdag.get(*basic_block)))) + StructuredBlock::Basic(Box::new(self.term_to_basic_block(basic_block))) }, ("Ite", [name, then_branch, else_branch]) => { - let string = Optimizer::string_term_to_string(&self.termdag.get(*name)); + let string = self.string_term_to_string(name); StructuredBlock::Ite( string.to_string(), - Box::new(self.term_to_structured_block(&self.termdag.get(*then_branch))), - Box::new(self.term_to_structured_block(&self.termdag.get(*else_branch))), + Box::new(self.term_to_structured_block(then_branch)), + Box::new(self.term_to_structured_block(else_branch)), ) }, ("Loop", [block]) => { - StructuredBlock::Loop(Box::new(self.term_to_structured_block(&self.termdag.get(*block)))) + StructuredBlock::Loop(Box::new(self.term_to_structured_block(block))) }, ("Sequence", [block, rest]) => StructuredBlock::Sequence(vec![ - self.term_to_structured_block(&self.termdag.get(*block)), - self.term_to_structured_block(&self.termdag.get(*rest)), + self.term_to_structured_block(block), + self.term_to_structured_block(rest), ]), ("Break", [n]) => { - if let Term::Lit(egglog::ast::Literal::Int(n)) = self.termdag.get(*n) { + if let Term::Lit(egglog::ast::Literal::Int(n)) = self.get(n) { StructuredBlock::Break(n.try_into().unwrap()) } else { panic!("expected int literal for break"); } }, ("Return", [val]) => { - match_term_app!(self.termdag.get(*val); { + match_term_app!(self.get(val); { ("Void", _) => StructuredBlock::Return(None), ("ReturnValue", [arg]) => { - match self.termdag.get(*arg) { + match self.get(arg) { Term::Lit(egglog::ast::Literal::String(s)) => { StructuredBlock::Return(Some(s.to_string())) } @@ -98,11 +102,11 @@ impl TermConverter<'_> { }) } - pub(crate) fn term_to_basic_block(&mut self, term: &Term) -> BasicBlock { - match_term_app!(term; { + pub(crate) fn term_to_basic_block(&mut self, id: &TermId) -> BasicBlock { + match_term_app!(self.get(id); { ("BlockNamed", [name, code]) => { - let name = Optimizer::string_term_to_string(&self.termdag.get(*name)); - let code_vec = self.term_conslist_to_vec(&self.termdag.get(*code), "Code"); + let name = self.string_term_to_string(name); + let code_vec = self.term_conslist_to_vec(code, "Code"); let mut instrs = vec![]; // let mut memo = HashMap::::new(); @@ -121,27 +125,27 @@ impl TermConverter<'_> { }) } - fn term_conslist_to_vec_helper(&self, term: &Term, res: &mut Vec, prefix: &str) { - match_term_app!(term; { + fn term_conslist_to_vec_helper(&self, id: &TermId, res: &mut Vec, prefix: &str) { + match_term_app!(self.get(id); { (op, [head, tail]) if op == prefix.to_string() + "Cons" => { - res.push(self.termdag.get(*head)); - self.term_conslist_to_vec_helper(&self.termdag.get(*tail), res, prefix); + res.push(*head); + self.term_conslist_to_vec_helper(tail, res, prefix); }, (op, []) if op == prefix.to_string() + "Nil" => {} (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) }) } - fn term_conslist_to_vec(&self, term: &Term, prefix: &str) -> Vec { + fn term_conslist_to_vec(&self, id: &TermId, prefix: &str) -> Vec { let mut res = vec![]; - self.term_conslist_to_vec_helper(term, &mut res, prefix); + self.term_conslist_to_vec_helper(id, &mut res, prefix); res } - fn term_to_instructions(&mut self, term: &Term, res: &mut Vec) { - match_term_app!(term; { + fn term_to_instructions(&mut self, id: &TermId, res: &mut Vec) { + match_term_app!(self.get(id); { ("Print", [arg]) => { - let arg = self.term_to_code(&self.termdag.get(*arg), res, None); + let arg = self.term_to_code(arg, res, None); res.push(Instruction::Effect { op: EffectOps::Print, @@ -153,13 +157,13 @@ impl TermConverter<'_> { }, ("End", []) => {}, ("Assign", [dest, src]) => { - let dest = Optimizer::string_term_to_string(&self.termdag.get(*dest)); - self.term_to_code(&self.termdag.get(*src), res, Some(dest.to_string())); + let dest = self.string_term_to_string(dest); + self.term_to_code(src, res, Some(dest.to_string())); }, (op @ ("store" | "free"), args) => { let args = args .iter() - .map(|arg| self.term_to_code(&self.termdag.get(*arg), res, None)) + .map(|arg| self.term_to_code(arg, res, None)) .collect::>(); @@ -172,9 +176,9 @@ impl TermConverter<'_> { }); }, ("alloc", [atype, dest, arg]) => { - let atype = self.term_to_type(&self.termdag.get(*atype)); - let dest = Optimizer::string_term_to_string(&self.termdag.get(*dest)); - let arg = self.term_to_code(&self.termdag.get(*arg), res, None); + let atype = self.term_to_type(atype); + let dest = self.string_term_to_string(dest); + let arg = self.term_to_code(arg, res, None); res.push(Instruction::Value { dest, args: vec![arg], @@ -191,7 +195,7 @@ impl TermConverter<'_> { pub(crate) fn term_to_code( &mut self, - term: &Term, + id: &TermId, res: &mut Vec, assign_to: Option, ) -> String { @@ -200,14 +204,14 @@ impl TermConverter<'_> { None => self.optimizer.fresh_var(), }; - match term { + match self.get(id) { Term::Lit(literal) => { res.push(Instruction::Constant { dest: dest.clone(), op: bril_rs::ConstOps::Const, - value: self.optimizer.literal_to_bril(literal), + value: self.optimizer.literal_to_bril(&literal), pos: None, - const_type: self.optimizer.literal_to_type(literal), + const_type: self.optimizer.literal_to_type(&literal), }); dest } @@ -218,31 +222,31 @@ impl TermConverter<'_> { var.to_string() } } - _ => { - match_term_app!(term; { + t => { + match_term_app!(t; { ("Var", [arg]) => { - match self.termdag.get(*arg) { + match self.get(arg) { Term::Lit(egglog::ast::Literal::String(var)) => var.to_string(), _ => panic!("expected string literal for var"), } }, - ("ReturnValue", [arg]) => self.term_to_code(&self.termdag.get(*arg), res, assign_to), + ("ReturnValue", [arg]) => self.term_to_code(arg, res, assign_to), (op @ ("True" | "False" | "Int" | "Float" | "Char"), [ty, args @ ..]) => { let lit = match (op, args) { ("True", []) => Literal::Bool(true), ("False", []) => Literal::Bool(false), ("Int", [arg]) => { - let arg = self.termdag.get(*arg); + let arg = self.get(arg); let arg_s = self.termdag.to_string(&arg); Literal::Int(arg_s.parse::().unwrap()) } ("Float", [arg]) => { - let arg = self.termdag.get(*arg); + let arg = self.get(arg); let arg_s = self.termdag.to_string(&arg); Literal::Float(arg_s.parse::().unwrap()) } ("Char", [arg]) => { - let arg = self.termdag.get(*arg); + let arg = self.get(arg); let arg_s = self.termdag.to_string(&arg); assert_eq!(arg_s.len(), 1); Literal::Char(arg_s.chars().next().unwrap()) @@ -254,16 +258,16 @@ impl TermConverter<'_> { op: bril_rs::ConstOps::Const, value: lit, pos: None, - const_type: self.term_to_type(&self.termdag.get(*ty)), + const_type: self.term_to_type(ty), }); dest }, ("phi", [etype, arg1, arg2, label1, label2]) => { - let etype = self.term_to_type(&self.termdag.get(*etype)); - let arg1 = self.term_to_code(&self.termdag.get(*arg1), res, None); - let arg2 = self.term_to_code(&self.termdag.get(*arg2), res, None); - let label1 = Optimizer::string_term_to_string(&self.termdag.get(*label1)); - let label2 = Optimizer::string_term_to_string(&self.termdag.get(*label2)); + let etype = self.term_to_type(etype); + let arg1 = self.term_to_code(arg1, res, None); + let arg2 = self.term_to_code(arg2, res, None); + let label1 = self.string_term_to_string(label1); + let label2 = self.string_term_to_string(label2); res.push(Instruction::Value { dest: dest.clone(), args: vec![arg1, arg2], @@ -277,11 +281,11 @@ impl TermConverter<'_> { }, (op, args) => { assert!(op != "Void"); - let etype = self.term_to_type(&self.termdag.get(args[0])); + let etype = self.term_to_type(&args[0]); let args_vars = args .iter() .skip(1) - .map(|arg| self.term_to_code(&self.termdag.get(*arg), res, None)) + .map(|arg| self.term_to_code(arg, res, None)) .collect::>(); res.push(Instruction::Value { dest: dest.clone(), @@ -299,16 +303,24 @@ impl TermConverter<'_> { } } - pub(crate) fn term_to_type(&self, term: &Term) -> Type { - match_term_app!(term; { + pub(crate) fn term_to_type(&self, id: &TermId) -> Type { + match_term_app!(self.get(id); { ("IntT", []) => Type::Int, ("BoolT", []) => Type::Bool, ("FloatT", []) => Type::Float, ("CharT", []) => Type::Char, - ("PointerT", [child]) => Type::Pointer(Box::new(self.term_to_type(&self.termdag.get(*child)))), + ("PointerT", [child]) => Type::Pointer(Box::new(self.term_to_type(child))), (head, _) => panic!("unexpected head {}, in {}:{}:{}", head, file!(), line!(), column!()) }) } + + fn string_term_to_string(&self, id: &TermId) -> String { + if let Term::Lit(egglog::ast::Literal::String(string)) = self.get(id) { + string.to_string() + } else { + panic!("expected string literal"); + } + } } impl Optimizer { @@ -321,7 +333,7 @@ impl Optimizer { optimizer: self, termdag, }; - converter.term_to_structured_func(term) + converter.term_to_structured_func(&termdag.lookup(term)) } pub(crate) fn func_to_expr(&mut self, func: &StructuredFunction) -> Expr { @@ -417,14 +429,6 @@ impl Optimizer { Expr::Call("Var".into(), vec![self.string_to_expr(string)]) } - fn string_term_to_string(term: &Term) -> String { - if let Term::Lit(egglog::ast::Literal::String(string)) = term { - string.to_string() - } else { - panic!("expected string literal"); - } - } - pub(crate) fn convert_basic_block(&mut self, block: &BasicBlock) -> Expr { // leave prints in order // leave any effects in order From b4ada72e79a9d97c88fe41e63b229b2b38b5eddb Mon Sep 17 00:00:00 2001 From: "James R. Wilcox" Date: Wed, 23 Aug 2023 17:20:33 -0700 Subject: [PATCH 5/5] Memoize expressions in basic blocks --- src/conversions.rs | 46 +++++++++++-------- tests/snapshots/files__add_naiive.snap | 3 +- tests/snapshots/files__add_naiive_no_opt.snap | 9 +--- .../files__block-diamond_naiive.snap | 4 +- .../files__block-diamond_naiive_no_opt.snap | 4 +- tests/snapshots/files__diamond_naiive.snap | 4 +- .../files__diamond_naiive_no_opt.snap | 4 +- .../files__two_fns_naiive_no_opt.snap | 8 +--- 8 files changed, 35 insertions(+), 47 deletions(-) diff --git a/src/conversions.rs b/src/conversions.rs index edc7c5854..9b55ff3c7 100644 --- a/src/conversions.rs +++ b/src/conversions.rs @@ -108,10 +108,10 @@ impl TermConverter<'_> { let name = self.string_term_to_string(name); let code_vec = self.term_conslist_to_vec(code, "Code"); let mut instrs = vec![]; - // let mut memo = HashMap::::new(); + let mut memo = HashMap::::new(); for t in code_vec { - self.term_to_instructions(&t, &mut instrs); + self.term_to_instructions(&t, &mut instrs, &mut memo); } BasicBlock { @@ -142,10 +142,15 @@ impl TermConverter<'_> { res } - fn term_to_instructions(&mut self, id: &TermId, res: &mut Vec) { + fn term_to_instructions( + &mut self, + id: &TermId, + res: &mut Vec, + memo: &mut HashMap, + ) { match_term_app!(self.get(id); { ("Print", [arg]) => { - let arg = self.term_to_code(arg, res, None); + let arg = self.term_to_code(arg, res, None, memo); res.push(Instruction::Effect { op: EffectOps::Print, @@ -158,12 +163,12 @@ impl TermConverter<'_> { ("End", []) => {}, ("Assign", [dest, src]) => { let dest = self.string_term_to_string(dest); - self.term_to_code(src, res, Some(dest.to_string())); + self.term_to_code(src, res, Some(dest.to_string()), memo); }, (op @ ("store" | "free"), args) => { let args = args .iter() - .map(|arg| self.term_to_code(arg, res, None)) + .map(|arg| self.term_to_code(arg, res, None, memo)) .collect::>(); @@ -178,7 +183,7 @@ impl TermConverter<'_> { ("alloc", [atype, dest, arg]) => { let atype = self.term_to_type(atype); let dest = self.string_term_to_string(dest); - let arg = self.term_to_code(arg, res, None); + let arg = self.term_to_code(arg, res, None, memo); res.push(Instruction::Value { dest, args: vec![arg], @@ -198,13 +203,18 @@ impl TermConverter<'_> { id: &TermId, res: &mut Vec, assign_to: Option, + memo: &mut HashMap, ) -> String { + if memo.contains_key(id) && assign_to.is_none() { + return memo[id].clone(); + } + let dest = match &assign_to { Some(dest) => dest.clone(), None => self.optimizer.fresh_var(), }; - match self.get(id) { + let ret = match self.get(id) { Term::Lit(literal) => { res.push(Instruction::Constant { dest: dest.clone(), @@ -215,13 +225,6 @@ impl TermConverter<'_> { }); dest } - Term::Var(var) => { - if let Some(_output) = assign_to { - panic!("Cannot assign var to var") - } else { - var.to_string() - } - } t => { match_term_app!(t; { ("Var", [arg]) => { @@ -230,7 +233,7 @@ impl TermConverter<'_> { _ => panic!("expected string literal for var"), } }, - ("ReturnValue", [arg]) => self.term_to_code(arg, res, assign_to), + ("ReturnValue", [arg]) => self.term_to_code(arg, res, assign_to, memo), (op @ ("True" | "False" | "Int" | "Float" | "Char"), [ty, args @ ..]) => { let lit = match (op, args) { ("True", []) => Literal::Bool(true), @@ -264,8 +267,8 @@ impl TermConverter<'_> { }, ("phi", [etype, arg1, arg2, label1, label2]) => { let etype = self.term_to_type(etype); - let arg1 = self.term_to_code(arg1, res, None); - let arg2 = self.term_to_code(arg2, res, None); + let arg1 = self.term_to_code(arg1, res, None, memo); + let arg2 = self.term_to_code(arg2, res, None, memo); let label1 = self.string_term_to_string(label1); let label2 = self.string_term_to_string(label2); res.push(Instruction::Value { @@ -285,7 +288,7 @@ impl TermConverter<'_> { let args_vars = args .iter() .skip(1) - .map(|arg| self.term_to_code(arg, res, None)) + .map(|arg| self.term_to_code(arg, res, None, memo)) .collect::>(); res.push(Instruction::Value { dest: dest.clone(), @@ -300,7 +303,10 @@ impl TermConverter<'_> { } }) } - } + }; + + memo.insert(*id, ret.clone()); + ret } pub(crate) fn term_to_type(&self, id: &TermId) -> Type { diff --git a/tests/snapshots/files__add_naiive.snap b/tests/snapshots/files__add_naiive.snap index a3e74871d..6543f4067 100644 --- a/tests/snapshots/files__add_naiive.snap +++ b/tests/snapshots/files__add_naiive.snap @@ -7,8 +7,7 @@ expression: "format!(\"{}\", res)" v0: int = const 1; v1: int = const 2; v2: int = const 3; - v0_: int = const 3; - print v0_; + print v2; ret; .sblock___0: .exit___: diff --git a/tests/snapshots/files__add_naiive_no_opt.snap b/tests/snapshots/files__add_naiive_no_opt.snap index 4932bfc19..f000bd986 100644 --- a/tests/snapshots/files__add_naiive_no_opt.snap +++ b/tests/snapshots/files__add_naiive_no_opt.snap @@ -6,13 +6,8 @@ expression: "format!(\"{}\", res)" .entry___: v0: int = const 1; v1: int = const 2; - v0_: int = const 1; - v1_: int = const 2; - v2: int = add v0_ v1_; - v3_: int = const 1; - v4_: int = const 2; - v2_: int = add v3_ v4_; - print v2_; + v2: int = add v0 v1; + print v2; ret; .sblock___0: .exit___: diff --git a/tests/snapshots/files__block-diamond_naiive.snap b/tests/snapshots/files__block-diamond_naiive.snap index 98de39b23..ae6f858d5 100644 --- a/tests/snapshots/files__block-diamond_naiive.snap +++ b/tests/snapshots/files__block-diamond_naiive.snap @@ -7,9 +7,7 @@ expression: "format!(\"{}\", res)" one: int = const 1; two: int = const 2; x: int = const 0; - v0_: int = const 1; - v1_: int = const 2; - a_cond: bool = lt v0_ v1_; + a_cond: bool = lt one two; br a_cond .sblock___4 .sblock___5; .sblock___4: jmp .sblock___2; diff --git a/tests/snapshots/files__block-diamond_naiive_no_opt.snap b/tests/snapshots/files__block-diamond_naiive_no_opt.snap index 98de39b23..ae6f858d5 100644 --- a/tests/snapshots/files__block-diamond_naiive_no_opt.snap +++ b/tests/snapshots/files__block-diamond_naiive_no_opt.snap @@ -7,9 +7,7 @@ expression: "format!(\"{}\", res)" one: int = const 1; two: int = const 2; x: int = const 0; - v0_: int = const 1; - v1_: int = const 2; - a_cond: bool = lt v0_ v1_; + a_cond: bool = lt one two; br a_cond .sblock___4 .sblock___5; .sblock___4: jmp .sblock___2; diff --git a/tests/snapshots/files__diamond_naiive.snap b/tests/snapshots/files__diamond_naiive.snap index 0514fdb8b..b5d444ac6 100644 --- a/tests/snapshots/files__diamond_naiive.snap +++ b/tests/snapshots/files__diamond_naiive.snap @@ -5,9 +5,7 @@ expression: "format!(\"{}\", res)" @main { .entry___: x: int = const 4; - v0_: int = const 4; - v1_: int = const 4; - cond: bool = lt v0_ v1_; + cond: bool = lt x x; br cond .sblock___3 .sblock___4; .sblock___3: jmp .sblock___1; diff --git a/tests/snapshots/files__diamond_naiive_no_opt.snap b/tests/snapshots/files__diamond_naiive_no_opt.snap index 0514fdb8b..b5d444ac6 100644 --- a/tests/snapshots/files__diamond_naiive_no_opt.snap +++ b/tests/snapshots/files__diamond_naiive_no_opt.snap @@ -5,9 +5,7 @@ expression: "format!(\"{}\", res)" @main { .entry___: x: int = const 4; - v0_: int = const 4; - v1_: int = const 4; - cond: bool = lt v0_ v1_; + cond: bool = lt x x; br cond .sblock___3 .sblock___4; .sblock___3: jmp .sblock___1; diff --git a/tests/snapshots/files__two_fns_naiive_no_opt.snap b/tests/snapshots/files__two_fns_naiive_no_opt.snap index 013fb98db..13fc24f68 100644 --- a/tests/snapshots/files__two_fns_naiive_no_opt.snap +++ b/tests/snapshots/files__two_fns_naiive_no_opt.snap @@ -6,9 +6,7 @@ expression: "format!(\"{}\", res)" .entry___: v0: int = const 1; v1: int = const 2; - v0_: int = const 1; - v1_: int = const 2; - v2: int = add v0_ v1_; + v2: int = add v0 v1; ret; .sblock___0: .exit___: @@ -17,9 +15,7 @@ expression: "format!(\"{}\", res)" .entry___: v0: int = const 1; v1: int = const 2; - v2_: int = const 1; - v3_: int = const 2; - v2: int = sub v2_ v3_; + v2: int = sub v0 v1; ret; .sblock___0: .exit___: