Skip to content

Commit

Permalink
Refactor atomize
Browse files Browse the repository at this point in the history
  • Loading branch information
@JonathanBrouwer
JonathanBrouwer committed Nov 25, 2023
1 parent e7049ad commit 17ee16e
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 109 deletions.
152 changes: 59 additions & 93 deletions compiler/src/passes/atomize/atomize.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,47 +36,33 @@ fn atomize_def(def: DefRevealed) -> DefAtomized {
}
}

fn atomize_expr<'p>(expr: Meta<Type<UniqueSym<'p>>, RExpr<'p>>) -> Meta<Type<UniqueSym<'p>>, AExpr<'p>> {
fn atomize_expr<'p>(
expr: Meta<Type<UniqueSym<'p>>, RExpr<'p>>,
) -> Meta<Type<UniqueSym<'p>>, AExpr<'p>> {
// Keep track of all the priors. These are bindings that should come before evaluating the expression.
let mut priors = Vec::new();

// Returns the atomized expression. This can hold `Var`s that must be bound by the priors.
let inner = match expr.inner {
RExpr::Lit { val } => AExpr::Atom {
atm: Atom::Val { val },
},
RExpr::Var { sym} => AExpr::Atom {
RExpr::Var { sym } => AExpr::Atom {
atm: Atom::Var { sym },
},
RExpr::BinaryOp { op, exprs: [lhs, rhs] } => {
let (lhs_arg, lhs_extra) = atomize_atom(*lhs);
let (rhs_arg, rhs_extra) = atomize_atom(*rhs);

[lhs_extra, rhs_extra]
.into_iter()
.flatten()
.rfold(AExpr::BinaryOp { op, exprs: [lhs_arg, rhs_arg] }, |bdy, (sym, bnd)| {
AExpr::Let {
sym,
bnd: Box::new(bnd),
bdy: Box::new(Meta {
inner: bdy,
meta: expr.meta.clone()
}),
}
})
RExpr::BinaryOp {
op,
exprs: [lhs, rhs],
} => AExpr::BinaryOp {
op,
exprs: [
atomize_atom(*lhs, &mut priors),
atomize_atom(*rhs, &mut priors),
],
},
RExpr::UnaryOp { op, expr: arg } => {
let (arg, extra) = atomize_atom(*arg);

extra
.into_iter()
.rfold(AExpr::UnaryOp { op, expr: arg }, |bdy, (sym, bnd)| {
AExpr::Let {
sym,
bnd: Box::new(bnd),
bdy: Box::new(Meta {
inner: bdy,
meta: expr.meta.clone()
}),
}
})
RExpr::UnaryOp { op, expr: arg } => AExpr::UnaryOp {
op,
expr: atomize_atom(*arg, &mut priors),
},
RExpr::Let { sym, bnd, bdy } => AExpr::Let {
sym,
Expand All @@ -93,28 +79,14 @@ fn atomize_expr<'p>(expr: Meta<Type<UniqueSym<'p>>, RExpr<'p>>) -> Meta<Type<Uni
unreachable!()
};

let (args, extras): (Vec<_>, Vec<_>) = args
.into_iter()
.map(atomize_atom)
.zip(params)
.map(|((arg, extra), arg_typ)| ((arg, arg_typ), extra))
.unzip();

let (fun, fun_expr) = atomize_atom(*fun);

fun_expr
.into_iter()
.chain(extras.into_iter().flatten())
.rfold(AExpr::Apply { fun, args}, |bdy, (sym, bnd)| {
AExpr::Let {
sym,
bnd: Box::new(bnd),
bdy: Box::new(Meta {
inner: bdy,
meta: expr.meta.clone()
}),
}
})
AExpr::Apply {
fun: atomize_atom(*fun, &mut priors),
args: args
.into_iter()
.map(|arg| atomize_atom(arg, &mut priors))
.zip(params)
.collect(),
}
}
RExpr::FunRef { sym } => AExpr::FunRef { sym },
RExpr::Loop { bdy } => AExpr::Loop {
Expand All @@ -135,55 +107,49 @@ fn atomize_expr<'p>(expr: Meta<Type<UniqueSym<'p>>, RExpr<'p>>) -> Meta<Type<Uni
RExpr::Return { bdy } => AExpr::Return {
bdy: Box::new(atomize_expr(*bdy)),
},
RExpr::Struct { sym, fields } => {
let (fields, extras): (Vec<_>, Vec<_>) = fields
RExpr::Struct { sym, fields } => AExpr::Struct {
sym,
fields: fields
.into_iter()
.map(|(sym, expr)| {
let (field, extra) = atomize_atom(expr);
((sym, field), extra)
let field = atomize_atom(expr, &mut priors);
(sym, field)
})
.unzip();

extras.into_iter().flatten().rfold(
AExpr::Struct { sym, fields },
|bdy , (sym, bnd)| AExpr::Let {
sym,
bnd: Box::new(bnd),
bdy: Box::new(Meta {
inner: bdy,
meta: expr.meta.clone()
}),
},
)
}
RExpr::AccessField { strct, field } => {
let (strct, extra) = atomize_atom(*strct);

extra.into_iter().rfold(
AExpr::AccessField { strct, field },
|bdy, (sym, bnd)| AExpr::Let {
sym,
bnd: Box::new(bnd),
bdy: Box::new(Meta {
inner: bdy,
meta: expr.meta.clone()
}),
},
)
}
.collect(),
},
RExpr::AccessField { strct, field } => AExpr::AccessField {
strct: atomize_atom(*strct, &mut priors),
field,
},
};

// Chains all the priors with the atomized expression as the body.
let inner = priors
.into_iter()
.rfold(inner, |bdy, (sym, bnd)| AExpr::Let {
sym,
bnd: Box::new(bnd),
bdy: Box::new(Meta {
inner: bdy,
meta: expr.meta.clone(),
}),
});

Meta {
inner,
meta: expr.meta
meta: expr.meta,
}
}

fn atomize_atom<'p>(expr: Meta<Type<UniqueSym<'p>>, RExpr<'p>>) -> (Atom<'p>, Option<(UniqueSym<'p>, Meta<Type<UniqueSym<'p>>, AExpr<'p>>)>) {
fn atomize_atom<'p>(
expr: Meta<Type<UniqueSym<'p>>, RExpr<'p>>,
priors: &mut Vec<(UniqueSym<'p>, Meta<Type<UniqueSym<'p>>, AExpr<'p>>)>,
) -> Atom<'p> {
if let RExpr::Lit { val } = expr.inner {
(Atom::Val { val }, None)
Atom::Val { val }
} else {
let tmp = gen_sym("tmp");
(Atom::Var { sym: tmp }, Some((tmp, atomize_expr(expr))))
priors.push((tmp, atomize_expr(expr)));
Atom::Var { sym: tmp }
}
}
25 changes: 12 additions & 13 deletions compiler/src/passes/atomize/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@ mod tests;

use crate::passes::parse::types::Type;
use crate::passes::parse::{BinaryOp, Def, Meta, UnaryOp};
use crate::passes::select::io::Std;
use crate::passes::validate::{DefValidated, ExprValidated, PrgValidated, TLit};
use crate::utils::gen_sym::UniqueSym;
use std::collections::HashMap;
use crate::passes::select::io::Std;

pub struct PrgAtomized<'p> {
pub defs: HashMap<UniqueSym<'p>, DefAtomized<'p>>,
Expand Down Expand Up @@ -123,7 +123,9 @@ impl<'p> From<DefAtomized<'p>> for DefValidated<'p> {
}
}

impl<'p> From<Meta<Type<UniqueSym<'p>>, AExpr<'p>>> for Meta<Type<UniqueSym<'p>>, ExprValidated<'p>> {
impl<'p> From<Meta<Type<UniqueSym<'p>>, AExpr<'p>>>
for Meta<Type<UniqueSym<'p>>, ExprValidated<'p>>
{
fn from(value: Meta<Type<UniqueSym<'p>>, AExpr<'p>>) -> Self {
let inner = match value.inner {
AExpr::Atom { atm, .. } => return atm.into(),
Expand Down Expand Up @@ -176,32 +178,29 @@ impl<'p> From<Meta<Type<UniqueSym<'p>>, AExpr<'p>>> for Meta<Type<UniqueSym<'p>>
.into_iter()
.map(|(sym, atm)| (sym, atm.into()))
.collect(),

},
AExpr::AccessField { strct, field } => ExprValidated::AccessField {
strct: Box::new(strct.into()),
field,

},
};

Meta {
inner,
meta: value.meta
meta: value.meta,
}
}
}

// Note that casting to Never here is safe because this `From` is only used by the interpreter which doesn't care about the type information.
impl<'p> From<Atom<'p>> for Meta<Type<UniqueSym<'p>>, ExprValidated<'p>> {
fn from(value: Atom<'p>) -> Self {
Meta { meta: Type::Never, inner: match value {
Atom::Val { val } => ExprValidated::Lit {
val,
},
Atom::Var { sym } => ExprValidated::Var {
sym,
Meta {
meta: Type::Never,
inner: match value {
Atom::Val { val } => ExprValidated::Lit { val },
Atom::Var { sym } => ExprValidated::Var { sym },
},
}}
}
}
}
7 changes: 4 additions & 3 deletions compiler/src/passes/atomize/tests.rs
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
use test_each_file::test_each_file;
use crate::interpreter::TestIO;
use crate::utils::split_test::split_test;
use crate::passes::parse::parse::parse_program;
use crate::passes::validate::PrgValidated;
use crate::utils::split_test::split_test;
use test_each_file::test_each_file;

fn atomize([test]: [&str; 1]) {
let (input, expected_output, expected_return, _) = split_test(test);

let program: PrgValidated = parse_program(test).unwrap()
let program: PrgValidated = parse_program(test)
.unwrap()
.validate()
.unwrap()
.reveal()
Expand Down
Empty file removed duplicate_arg
View file
Empty file.

0 comments on commit 17ee16e

Please sign in to comment.