Add AST optimizer.

This commit is contained in:
Stephen Chung 2020-03-09 21:57:07 +08:00
parent c467ffd58d
commit 55e7af7b04
5 changed files with 349 additions and 88 deletions

View File

@ -95,13 +95,13 @@ impl<'e> Engine<'e> {
} }
/// Compile a string into an AST /// Compile a string into an AST
pub fn compile(input: &str) -> Result<AST, ParseError> { pub fn compile(&self, input: &str) -> Result<AST, ParseError> {
let tokens = lex(input); let tokens = lex(input);
parse(&mut tokens.peekable()) parse(&mut tokens.peekable(), self.optimize)
} }
/// Compile a file into an AST /// Compile a file into an AST
pub fn compile_file(filename: &str) -> Result<AST, EvalAltResult> { pub fn compile_file(&self, filename: &str) -> Result<AST, EvalAltResult> {
use std::fs::File; use std::fs::File;
use std::io::prelude::*; use std::io::prelude::*;
@ -112,7 +112,7 @@ impl<'e> Engine<'e> {
f.read_to_string(&mut contents) f.read_to_string(&mut contents)
.map_err(|err| EvalAltResult::ErrorReadingScriptFile(filename.into(), err)) .map_err(|err| EvalAltResult::ErrorReadingScriptFile(filename.into(), err))
.and_then(|_| Self::compile(&contents).map_err(EvalAltResult::ErrorParsing)) .and_then(|_| self.compile(&contents).map_err(EvalAltResult::ErrorParsing))
} }
/// Evaluate a file /// Evaluate a file
@ -142,7 +142,7 @@ impl<'e> Engine<'e> {
scope: &mut Scope, scope: &mut Scope,
input: &str, input: &str,
) -> Result<T, EvalAltResult> { ) -> Result<T, EvalAltResult> {
let ast = Self::compile(input).map_err(EvalAltResult::ErrorParsing)?; let ast = self.compile(input).map_err(EvalAltResult::ErrorParsing)?;
self.eval_ast_with_scope(scope, &ast) self.eval_ast_with_scope(scope, &ast)
} }
@ -229,7 +229,7 @@ impl<'e> Engine<'e> {
) -> Result<(), EvalAltResult> { ) -> Result<(), EvalAltResult> {
let tokens = lex(input); let tokens = lex(input);
parse(&mut tokens.peekable()) parse(&mut tokens.peekable(), self.optimize)
.map_err(|err| EvalAltResult::ErrorParsing(err)) .map_err(|err| EvalAltResult::ErrorParsing(err))
.and_then(|AST(ref statements, ref functions)| { .and_then(|AST(ref statements, ref functions)| {
for f in functions { for f in functions {
@ -258,11 +258,12 @@ impl<'e> Engine<'e> {
/// # Example /// # Example
/// ///
/// ```rust /// ```rust
/// # use rhai::{Engine, EvalAltResult}; /// # fn main() -> Result<(), rhai::EvalAltResult> {
/// # fn main() -> Result<(), EvalAltResult> { /// use rhai::Engine;
///
/// let mut engine = Engine::new(); /// let mut engine = Engine::new();
/// ///
/// let ast = Engine::compile("fn add(x, y) { x.len() + y }")?; /// let ast = engine.compile("fn add(x, y) { x.len() + y }")?;
/// ///
/// let result: i64 = engine.call_fn("add", &ast, (&mut String::from("abc"), &mut 123_i64))?; /// let result: i64 = engine.call_fn("add", &ast, (&mut String::from("abc"), &mut 123_i64))?;
/// ///
@ -309,16 +310,20 @@ impl<'e> Engine<'e> {
/// # Example /// # Example
/// ///
/// ```rust /// ```rust
/// # use rhai::Engine; /// # fn main() -> Result<(), rhai::EvalAltResult> {
/// use rhai::Engine;
///
/// let mut result = String::from(""); /// let mut result = String::from("");
/// { /// {
/// let mut engine = Engine::new(); /// let mut engine = Engine::new();
/// ///
/// // Override action of 'print' function /// // Override action of 'print' function
/// engine.on_print(|s| result.push_str(s)); /// engine.on_print(|s| result.push_str(s));
/// engine.consume("print(40 + 2);").unwrap(); /// engine.consume("print(40 + 2);")?;
/// } /// }
/// assert_eq!(result, "42"); /// assert_eq!(result, "42");
/// # Ok(())
/// # }
/// ``` /// ```
pub fn on_print(&mut self, callback: impl FnMut(&str) + 'e) { pub fn on_print(&mut self, callback: impl FnMut(&str) + 'e) {
self.on_print = Box::new(callback); self.on_print = Box::new(callback);
@ -329,16 +334,20 @@ impl<'e> Engine<'e> {
/// # Example /// # Example
/// ///
/// ```rust /// ```rust
/// # use rhai::Engine; /// # fn main() -> Result<(), rhai::EvalAltResult> {
/// use rhai::Engine;
///
/// let mut result = String::from(""); /// let mut result = String::from("");
/// { /// {
/// let mut engine = Engine::new(); /// let mut engine = Engine::new();
/// ///
/// // Override action of 'debug' function /// // Override action of 'debug' function
/// engine.on_debug(|s| result.push_str(s)); /// engine.on_debug(|s| result.push_str(s));
/// engine.consume(r#"debug("hello");"#).unwrap(); /// engine.consume(r#"debug("hello");"#)?;
/// } /// }
/// assert_eq!(result, "\"hello\""); /// assert_eq!(result, "\"hello\"");
/// # Ok(())
/// # }
/// ``` /// ```
pub fn on_debug(&mut self, callback: impl FnMut(&str) + 'e) { pub fn on_debug(&mut self, callback: impl FnMut(&str) + 'e) {
self.on_debug = Box::new(callback); self.on_debug = Box::new(callback);

View File

@ -4,7 +4,7 @@ use crate::any::{Any, AnyExt, Dynamic, Variant};
use crate::parser::{Expr, FnDef, Position, Stmt}; use crate::parser::{Expr, FnDef, Position, Stmt};
use crate::result::EvalAltResult; use crate::result::EvalAltResult;
use crate::scope::Scope; use crate::scope::Scope;
use std::any::TypeId; use std::any::{type_name, TypeId};
use std::borrow::Cow; use std::borrow::Cow;
use std::cmp::{PartialEq, PartialOrd}; use std::cmp::{PartialEq, PartialOrd};
use std::collections::HashMap; use std::collections::HashMap;
@ -20,11 +20,11 @@ pub type FnAny = dyn Fn(FnCallArgs, Position) -> Result<Dynamic, EvalAltResult>;
type IteratorFn = dyn Fn(&Dynamic) -> Box<dyn Iterator<Item = Dynamic>>; type IteratorFn = dyn Fn(&Dynamic) -> Box<dyn Iterator<Item = Dynamic>>;
const KEYWORD_PRINT: &'static str = "print"; pub(crate) const KEYWORD_PRINT: &'static str = "print";
const KEYWORD_DEBUG: &'static str = "debug"; pub(crate) const KEYWORD_DEBUG: &'static str = "debug";
const KEYWORD_TYPE_OF: &'static str = "type_of"; pub(crate) const KEYWORD_TYPE_OF: &'static str = "type_of";
const FUNC_GETTER: &'static str = "get$"; pub(crate) const FUNC_GETTER: &'static str = "get$";
const FUNC_SETTER: &'static str = "set$"; pub(crate) const FUNC_SETTER: &'static str = "set$";
#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)] #[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]
enum IndexSourceType { enum IndexSourceType {
@ -42,17 +42,20 @@ pub struct FnSpec<'a> {
/// Rhai main scripting engine. /// Rhai main scripting engine.
/// ///
/// ```rust /// ```rust
/// # fn main() -> Result<(), rhai::EvalAltResult> {
/// use rhai::Engine; /// use rhai::Engine;
/// ///
/// fn main() { /// let mut engine = Engine::new();
/// let mut engine = Engine::new();
/// ///
/// if let Ok(result) = engine.eval::<i64>("40 + 2") { /// let result = engine.eval::<i64>("40 + 2")?;
/// println!("Answer: {}", result); // prints 42 ///
/// } /// println!("Answer: {}", result); // prints 42
/// } /// # Ok(())
/// # }
/// ``` /// ```
pub struct Engine<'e> { pub struct Engine<'e> {
/// Optimize the AST after compilation
pub(crate) optimize: bool,
/// A hashmap containing all compiled functions known to the engine /// A hashmap containing all compiled functions known to the engine
pub(crate) ext_functions: HashMap<FnSpec<'e>, Arc<FnIntExt<'e>>>, pub(crate) ext_functions: HashMap<FnSpec<'e>, Arc<FnIntExt<'e>>>,
/// A hashmap containing all script-defined functions /// A hashmap containing all script-defined functions
@ -72,6 +75,42 @@ pub enum FnIntExt<'a> {
} }
impl Engine<'_> { impl Engine<'_> {
/// Create a new `Engine`
pub fn new() -> Self {
// User-friendly names for built-in types
let type_names = [
(type_name::<String>(), "string"),
(type_name::<Array>(), "array"),
(type_name::<Dynamic>(), "dynamic"),
]
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect();
// Create the new scripting Engine
let mut engine = Engine {
optimize: true,
ext_functions: HashMap::new(),
script_functions: HashMap::new(),
type_iterators: HashMap::new(),
type_names,
on_print: Box::new(default_print), // default print/debug implementations
on_debug: Box::new(default_print),
};
engine.register_core_lib();
#[cfg(not(feature = "no_stdlib"))]
engine.register_stdlib(); // Register the standard library when no_stdlib is not set
engine
}
/// Control whether the `Engine` will optimize an AST after compilation
pub fn set_optimization(&mut self, optimize: bool) {
self.optimize = optimize
}
/// Universal method for calling functions, that are either /// Universal method for calling functions, that are either
/// registered with the `Engine` or written in Rhai /// registered with the `Engine` or written in Rhai
pub(crate) fn call_fn_raw( pub(crate) fn call_fn_raw(
@ -151,7 +190,7 @@ impl Engine<'_> {
); );
// Evaluate // Evaluate
match self.eval_stmt(&mut scope, &*func.body) { match self.eval_stmt(&mut scope, &func.body) {
// Convert return statement to return value // Convert return statement to return value
Err(EvalAltResult::Return(x, _)) => Ok(x), Err(EvalAltResult::Return(x, _)) => Ok(x),
other => other, other => other,
@ -726,7 +765,7 @@ impl Engine<'_> {
.map(|(_, _, _, x)| x), .map(|(_, _, _, x)| x),
// Statement block // Statement block
Expr::Block(block, _) => self.eval_stmt(scope, block), Expr::Stmt(stmt, _) => self.eval_stmt(scope, stmt),
// lhs = rhs // lhs = rhs
Expr::Assignment(lhs, rhs, _) => { Expr::Assignment(lhs, rhs, _) => {
@ -850,11 +889,14 @@ impl Engine<'_> {
stmt: &Stmt, stmt: &Stmt,
) -> Result<Dynamic, EvalAltResult> { ) -> Result<Dynamic, EvalAltResult> {
match stmt { match stmt {
// No-op
Stmt::Noop(_) => Ok(().into_dynamic()),
// Expression as statement // Expression as statement
Stmt::Expr(expr) => self.eval_expr(scope, expr), Stmt::Expr(expr) => self.eval_expr(scope, expr),
// Block scope // Block scope
Stmt::Block(block) => { Stmt::Block(block, _) => {
let prev_len = scope.len(); let prev_len = scope.len();
let mut last_result: Result<Dynamic, EvalAltResult> = Ok(().into_dynamic()); let mut last_result: Result<Dynamic, EvalAltResult> = Ok(().into_dynamic());
@ -988,38 +1030,6 @@ impl Engine<'_> {
.map(|s| s.as_str()) .map(|s| s.as_str())
.unwrap_or(name) .unwrap_or(name)
} }
/// Make a new engine
pub fn new<'a>() -> Engine<'a> {
use std::any::type_name;
// User-friendly names for built-in types
let type_names = [
(type_name::<String>(), "string"),
(type_name::<Array>(), "array"),
(type_name::<Dynamic>(), "dynamic"),
]
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect();
// Create the new scripting Engine
let mut engine = Engine {
ext_functions: HashMap::new(),
script_functions: HashMap::new(),
type_iterators: HashMap::new(),
type_names,
on_print: Box::new(default_print), // default print/debug implementations
on_debug: Box::new(default_print),
};
engine.register_core_lib();
#[cfg(not(feature = "no_stdlib"))]
engine.register_stdlib(); // Register the standard library when no_stdlib is not set
engine
}
} }
/// Print/debug to stdout /// Print/debug to stdout

View File

@ -68,7 +68,7 @@ mod call;
mod engine; mod engine;
mod error; mod error;
mod fn_register; mod fn_register;
//mod optimize; mod optimize;
mod parser; mod parser;
mod result; mod result;
mod scope; mod scope;

212
src/optimize.rs Normal file
View File

@ -0,0 +1,212 @@
use crate::parser::{Expr, Stmt};
fn optimize_stmt(stmt: Stmt, changed: &mut bool) -> Stmt {
match stmt {
Stmt::IfElse(expr, stmt1, None) => match *expr {
Expr::False(pos) => {
*changed = true;
Stmt::Noop(pos)
}
Expr::True(_) => optimize_stmt(*stmt1, changed),
_ => Stmt::IfElse(
Box::new(optimize_expr(*expr, changed)),
Box::new(optimize_stmt(*stmt1, changed)),
None,
),
},
Stmt::IfElse(expr, stmt1, Some(stmt2)) => match *expr {
Expr::False(_) => optimize_stmt(*stmt2, changed),
Expr::True(_) => optimize_stmt(*stmt1, changed),
_ => Stmt::IfElse(
Box::new(optimize_expr(*expr, changed)),
Box::new(optimize_stmt(*stmt1, changed)),
Some(Box::new(optimize_stmt(*stmt2, changed))),
),
},
Stmt::While(expr, stmt) => match *expr {
Expr::False(pos) => {
*changed = true;
Stmt::Noop(pos)
}
Expr::True(_) => Stmt::Loop(Box::new(optimize_stmt(*stmt, changed))),
_ => Stmt::While(
Box::new(optimize_expr(*expr, changed)),
Box::new(optimize_stmt(*stmt, changed)),
),
},
Stmt::Loop(stmt) => Stmt::Loop(Box::new(optimize_stmt(*stmt, changed))),
Stmt::For(id, expr, stmt) => Stmt::For(
id,
Box::new(optimize_expr(*expr, changed)),
Box::new(optimize_stmt(*stmt, changed)),
),
Stmt::Let(id, Some(expr), pos) => {
Stmt::Let(id, Some(Box::new(optimize_expr(*expr, changed))), pos)
}
Stmt::Let(_, None, _) => stmt,
Stmt::Block(statements, pos) => {
let original_len = statements.len();
let mut result: Vec<_> = statements
.into_iter() // For each statement
.map(|s| optimize_stmt(s, changed)) // Optimize the statement
.filter(Stmt::is_op) // Remove no-op's
.collect();
*changed = *changed || original_len != result.len();
match result[..] {
[] => {
// No statements in block - change to No-op
*changed = true;
Stmt::Noop(pos)
}
[Stmt::Let(_, _, _)] => {
// Only one let statement, but cannot promote
// (otherwise the variable gets declared in the scope above)
// and still need to run just in case there are side effects
Stmt::Block(result, pos)
}
[_] => {
// No statements in block - change to No-op
*changed = true;
result.remove(0)
}
_ => Stmt::Block(result, pos),
}
}
Stmt::Expr(expr) => Stmt::Expr(Box::new(optimize_expr(*expr, changed))),
Stmt::ReturnWithVal(Some(expr), is_return, pos) => Stmt::ReturnWithVal(
Some(Box::new(optimize_expr(*expr, changed))),
is_return,
pos,
),
Stmt::ReturnWithVal(None, _, _) => stmt,
Stmt::Noop(_) | Stmt::Break(_) => stmt,
}
}
fn optimize_expr(expr: Expr, changed: &mut bool) -> Expr {
match expr {
Expr::IntegerConstant(_, _)
| Expr::FloatConstant(_, _)
| Expr::Identifier(_, _)
| Expr::CharConstant(_, _)
| Expr::StringConstant(_, _)
| Expr::True(_)
| Expr::False(_)
| Expr::Unit(_) => expr,
Expr::Stmt(stmt, pos) => match optimize_stmt(*stmt, changed) {
Stmt::Noop(_) => {
*changed = true;
Expr::Unit(pos)
}
Stmt::Expr(expr) => {
*changed = true;
*expr
}
stmt => Expr::Stmt(Box::new(stmt), pos),
},
Expr::Assignment(id, expr, pos) => {
Expr::Assignment(id, Box::new(optimize_expr(*expr, changed)), pos)
}
Expr::Dot(lhs, rhs, pos) => Expr::Dot(
Box::new(optimize_expr(*lhs, changed)),
Box::new(optimize_expr(*rhs, changed)),
pos,
),
Expr::Index(lhs, rhs, pos) => Expr::Index(
Box::new(optimize_expr(*lhs, changed)),
Box::new(optimize_expr(*rhs, changed)),
pos,
),
Expr::Array(items, pos) => {
let original_len = items.len();
let items: Vec<_> = items
.into_iter()
.map(|expr| optimize_expr(expr, changed))
.collect();
*changed = *changed || original_len != items.len();
Expr::Array(items, pos)
}
Expr::And(lhs, rhs) => match (*lhs, *rhs) {
(Expr::True(_), rhs) => {
*changed = true;
rhs
}
(Expr::False(pos), _) => {
*changed = true;
Expr::False(pos)
}
(lhs, Expr::True(_)) => {
*changed = true;
lhs
}
(lhs, rhs) => Expr::And(
Box::new(optimize_expr(lhs, changed)),
Box::new(optimize_expr(rhs, changed)),
),
},
Expr::Or(lhs, rhs) => match (*lhs, *rhs) {
(Expr::False(_), rhs) => {
*changed = true;
rhs
}
(Expr::True(pos), _) => {
*changed = true;
Expr::True(pos)
}
(lhs, Expr::False(_)) => {
*changed = true;
lhs
}
(lhs, rhs) => Expr::Or(
Box::new(optimize_expr(lhs, changed)),
Box::new(optimize_expr(rhs, changed)),
),
},
Expr::FunctionCall(id, args, def_value, pos) => {
let original_len = args.len();
let args: Vec<_> = args
.into_iter()
.map(|a| optimize_expr(a, changed))
.collect();
*changed = *changed || original_len != args.len();
Expr::FunctionCall(id, args, def_value, pos)
}
}
}
pub(crate) fn optimize(mut statements: Vec<Stmt>) -> Vec<Stmt> {
loop {
let mut changed = false;
statements = statements
.into_iter()
.map(|stmt| optimize_stmt(stmt, &mut changed))
.filter(Stmt::is_op)
.collect();
if !changed {
break;
}
}
statements
}

View File

@ -2,7 +2,8 @@
use crate::any::Dynamic; use crate::any::Dynamic;
use crate::error::{LexError, ParseError, ParseErrorType}; use crate::error::{LexError, ParseError, ParseErrorType};
use std::{borrow::Cow, char, fmt, iter::Peekable, str::Chars, usize}; use crate::optimize::optimize;
use std::{borrow::Cow, char, fmt, iter::Peekable, str::Chars, str::FromStr, usize};
type LERR = LexError; type LERR = LexError;
type PERR = ParseErrorType; type PERR = ParseErrorType;
@ -128,23 +129,33 @@ pub struct AST(pub(crate) Vec<Stmt>, pub(crate) Vec<FnDef<'static>>);
pub struct FnDef<'a> { pub struct FnDef<'a> {
pub name: Cow<'a, str>, pub name: Cow<'a, str>,
pub params: Vec<Cow<'a, str>>, pub params: Vec<Cow<'a, str>>,
pub body: Box<Stmt>, pub body: Stmt,
pub pos: Position, pub pos: Position,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub enum Stmt { pub enum Stmt {
Noop(Position),
IfElse(Box<Expr>, Box<Stmt>, Option<Box<Stmt>>), IfElse(Box<Expr>, Box<Stmt>, Option<Box<Stmt>>),
While(Box<Expr>, Box<Stmt>), While(Box<Expr>, Box<Stmt>),
Loop(Box<Stmt>), Loop(Box<Stmt>),
For(String, Box<Expr>, Box<Stmt>), For(String, Box<Expr>, Box<Stmt>),
Let(String, Option<Box<Expr>>, Position), Let(String, Option<Box<Expr>>, Position),
Block(Vec<Stmt>), Block(Vec<Stmt>, Position),
Expr(Box<Expr>), Expr(Box<Expr>),
Break(Position), Break(Position),
ReturnWithVal(Option<Box<Expr>>, bool, Position), ReturnWithVal(Option<Box<Expr>>, bool, Position),
} }
impl Stmt {
pub fn is_op(&self) -> bool {
match self {
Stmt::Noop(_) => false,
_ => true,
}
}
}
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub enum Expr { pub enum Expr {
IntegerConstant(i64, Position), IntegerConstant(i64, Position),
@ -152,7 +163,7 @@ pub enum Expr {
Identifier(String, Position), Identifier(String, Position),
CharConstant(char, Position), CharConstant(char, Position),
StringConstant(String, Position), StringConstant(String, Position),
Block(Box<Stmt>, Position), Stmt(Box<Stmt>, Position),
FunctionCall(String, Vec<Expr>, Option<Dynamic>, Position), FunctionCall(String, Vec<Expr>, Option<Dynamic>, Position),
Assignment(Box<Expr>, Box<Expr>, Position), Assignment(Box<Expr>, Box<Expr>, Position),
Dot(Box<Expr>, Box<Expr>, Position), Dot(Box<Expr>, Box<Expr>, Position),
@ -174,7 +185,7 @@ impl Expr {
| Expr::CharConstant(_, pos) | Expr::CharConstant(_, pos)
| Expr::StringConstant(_, pos) | Expr::StringConstant(_, pos)
| Expr::FunctionCall(_, _, _, pos) | Expr::FunctionCall(_, _, _, pos)
| Expr::Block(_, pos) | Expr::Stmt(_, pos)
| Expr::Array(_, pos) | Expr::Array(_, pos)
| Expr::True(pos) | Expr::True(pos)
| Expr::False(pos) | Expr::False(pos)
@ -666,24 +677,23 @@ impl<'a> TokenIterator<'a> {
let out: String = result.iter().skip(2).filter(|&&c| c != '_').collect(); let out: String = result.iter().skip(2).filter(|&&c| c != '_').collect();
return Some(( return Some((
if let Ok(val) = i64::from_str_radix(&out, radix) { i64::from_str_radix(&out, radix)
Token::IntegerConstant(val) .map(Token::IntegerConstant)
} else { .unwrap_or_else(|_| {
Token::LexError(LERR::MalformedNumber(result.iter().collect())) Token::LexError(LERR::MalformedNumber(result.iter().collect()))
}, }),
pos, pos,
)); ));
} else { } else {
let out: String = result.iter().filter(|&&c| c != '_').collect(); let out: String = result.iter().filter(|&&c| c != '_').collect();
return Some(( return Some((
if let Ok(val) = out.parse::<i64>() { i64::from_str(&out)
Token::IntegerConstant(val) .map(Token::IntegerConstant)
} else if let Ok(val) = out.parse::<f64>() { .or_else(|_| f64::from_str(&out).map(Token::FloatConstant))
Token::FloatConstant(val) .unwrap_or_else(|_| {
} else { Token::LexError(LERR::MalformedNumber(result.iter().collect()))
Token::LexError(LERR::MalformedNumber(result.iter().collect())) }),
},
pos, pos,
)); ));
} }
@ -1288,7 +1298,7 @@ fn parse_primary<'a>(input: &mut Peekable<TokenIterator<'a>>) -> Result<Expr, Pa
// Block statement as expression // Block statement as expression
match input.peek() { match input.peek() {
Some(&(Token::LeftBrace, pos)) => { Some(&(Token::LeftBrace, pos)) => {
return parse_block(input).map(|block| Expr::Block(Box::new(block), pos)) return parse_block(input).map(|block| Expr::Stmt(Box::new(block), pos))
} }
_ => (), _ => (),
} }
@ -1667,7 +1677,7 @@ fn parse_block<'a>(input: &mut Peekable<TokenIterator<'a>>) -> Result<Stmt, Pars
None => return Err(ParseError::new(PERR::MissingLeftBrace, Position::eof())), None => return Err(ParseError::new(PERR::MissingLeftBrace, Position::eof())),
} }
input.next(); let pos = input.next().unwrap().1;
let mut statements = Vec::new(); let mut statements = Vec::new();
@ -1695,7 +1705,7 @@ fn parse_block<'a>(input: &mut Peekable<TokenIterator<'a>>) -> Result<Stmt, Pars
match input.peek() { match input.peek() {
Some(&(Token::RightBrace, _)) => { Some(&(Token::RightBrace, _)) => {
input.next(); input.next();
Ok(Stmt::Block(statements)) Ok(Stmt::Block(statements, pos))
} }
Some(&(_, pos)) => Err(ParseError::new( Some(&(_, pos)) => Err(ParseError::new(
PERR::MissingRightBrace("end of block".into()), PERR::MissingRightBrace("end of block".into()),
@ -1811,13 +1821,16 @@ fn parse_fn<'a>(input: &mut Peekable<TokenIterator<'a>>) -> Result<FnDef<'static
Ok(FnDef { Ok(FnDef {
name: name.into(), name: name.into(),
params: params, params,
body: Box::new(body), body,
pos: pos, pos,
}) })
} }
fn parse_top_level<'a>(input: &mut Peekable<TokenIterator<'a>>) -> Result<AST, ParseError> { fn parse_top_level<'a>(
input: &mut Peekable<TokenIterator<'a>>,
optimize_ast: bool,
) -> Result<AST, ParseError> {
let mut statements = Vec::new(); let mut statements = Vec::new();
let mut functions = Vec::new(); let mut functions = Vec::new();
@ -1833,9 +1846,26 @@ fn parse_top_level<'a>(input: &mut Peekable<TokenIterator<'a>>) -> Result<AST, P
} }
} }
Ok(AST(statements, functions)) return Ok(if optimize_ast {
AST(
optimize(statements),
functions
.into_iter()
.map(|mut fn_def| {
let mut body = optimize(vec![fn_def.body]);
fn_def.body = body.pop().unwrap();
fn_def
})
.collect(),
)
} else {
AST(statements, functions)
});
} }
pub fn parse<'a>(input: &mut Peekable<TokenIterator<'a>>) -> Result<AST, ParseError> { pub fn parse<'a>(
parse_top_level(input) input: &mut Peekable<TokenIterator<'a>>,
optimize_ast: bool,
) -> Result<AST, ParseError> {
parse_top_level(input, optimize_ast)
} }