schala/schala-lang/language/src/reduced_ir/mod.rs

use crate::ast;
use crate::symbol_table::{DefId, Symbol, SymbolSpec, SymbolTable};
use crate::builtin::Builtin;

use std::str::FromStr;
use std::collections::HashMap;
use std::rc::Rc;

mod test;

pub fn reduce(ast: &ast::AST, symbol_table: &SymbolTable) -> ReducedIR {
    let reducer = Reducer::new(symbol_table);
    reducer.reduce(ast)
}

struct Reducer<'a> {
    symbol_table: &'a SymbolTable,
    functions: HashMap<DefId, FunctionDefinition>,
}

impl<'a> Reducer<'a> {
    fn new(symbol_table: &'a SymbolTable) -> Self {
        Self {
            symbol_table,
            functions: HashMap::new(),
        }
    }

    fn reduce(mut self, ast: &ast::AST) -> ReducedIR {
        // First reduce all functions
        // TODO once this works, maybe rewrite it using the Visitor
        for statement in ast.statements.iter() {
            self.top_level_statement(&statement);
        }

        // Then compute the entrypoint statements (which may reference previously-computed
        // functions by ID)
        let mut entrypoint = vec![];
        for statement in ast.statements.iter() {
            let ast::Statement { id: item_id, kind, .. } = statement;
            match &kind {
                ast::StatementKind::Expression(expr) => {
                    entrypoint.push(Statement::Expression(self.expression(&expr)));
                },
                ast::StatementKind::Declaration(ast::Declaration::Binding { name, constant, expr, ..}) => {
                    let symbol = self.symbol_table.lookup_symbol(item_id).unwrap();
                    entrypoint.push(Statement::Binding { id: symbol.def_id.clone(), constant: *constant, expr: self.expression(&expr) });
                },
                _ => ()
            }
        }

        ReducedIR {
            functions: self.functions,
            entrypoint,
        }
    }

    fn top_level_statement(&mut self, statement: &ast::Statement) {
        let ast::Statement { id: item_id, kind, .. } = statement;
        match kind {
            ast::StatementKind::Expression(_expr) => {
                //TODO expressions can in principle contain definitions, but I won't worry
                //about it now
                ()
            },
            ast::StatementKind::Declaration(decl) => match decl {
                ast::Declaration::FuncDecl(_, statements) => {
                    self.insert_function_definition(item_id, statements);
                },
                _ => ()
            },
            ast::StatementKind::Import(..) => (),
            ast::StatementKind::Module(modspec) => {
                //TODO handle modules
                ()
            }
        }
    }

    fn function_internal_statement(&mut self, statement: &ast::Statement) -> Option<Statement> {
        let ast::Statement { id: item_id, kind, .. } = statement;
        match kind {
            ast::StatementKind::Expression(expr) => {
                Some(Statement::Expression(self.expression(expr)))
            },
            ast::StatementKind::Declaration(decl) => match decl {
                ast::Declaration::FuncDecl(_, statements) => {
                    self.insert_function_definition(item_id, statements);
                    None
                },
                _ => None
            },
            _ => None
        }
    }

    fn insert_function_definition(&mut self, item_id: &ast::ItemId, statements: &ast::Block) {
        let symbol = self.symbol_table.lookup_symbol(item_id).unwrap();
        let def_id = symbol.def_id.clone();
        let function_def = FunctionDefinition {
            body: self.function(statements)
        };
        self.functions.insert(def_id, function_def);
    }

    fn expression(&mut self, expr: &ast::Expression) -> Expression {
        use crate::ast::ExpressionKind::*;
        use Expression::{Unimplemented};

        match &expr.kind {
            NatLiteral(n) => Expression::Literal(Literal::Nat(*n)),
            FloatLiteral(f) => Expression::Literal(Literal::Float(*f)),
            StringLiteral(s) => Expression::Literal(Literal::StringLit(s.clone())),
            BoolLiteral(b) => Expression::Literal(Literal::Bool(*b)),
            BinExp(binop, lhs, rhs) =>  self.binop(binop, lhs, rhs),
            PrefixExp(op, arg) => self.prefix(op, arg),
            Value(qualified_name) => self.value(qualified_name),
            Call { f, arguments } => Unimplemented, // self.reduce_call_expression(f, arguments),
            TupleLiteral(exprs) => Expression::Tuple(exprs.iter().map(|e| self.expression(e)).collect()),
            IfExpression {
                discriminator,
                body,
            } => Unimplemented, //self.reduce_if_expression(deref_optional_box(discriminator), body),
            Lambda { params, body, .. } => Unimplemented, //self.reduce_lambda(params, body),
            NamedStruct { name, fields } => Unimplemented, //self.reduce_named_struct(name, fields),
            Index { .. } => Unimplemented,
            WhileExpression { .. } => Unimplemented,
            ForExpression { .. } => Unimplemented,
            ListLiteral { .. } => Unimplemented,
        }
    }

    fn function(&mut self, statements: &ast::Block) -> Vec<Statement> {
        statements.iter().filter_map(|stmt| self.function_internal_statement(stmt)).collect()
    }

    fn prefix(&mut self, prefix: &ast::PrefixOp, arg: &ast::Expression) -> Expression {
        let builtin: Option<Builtin> = TryFrom::try_from(prefix).ok();
        match builtin {
            Some(op) => {
                Expression::Call {
                    f: Box::new(Expression::Callable(Function::Builtin(op))),
                    args: vec![self.expression(arg)],
                }
            }
            None => {
                //TODO need this for custom prefix ops
                Expression::Unimplemented
            }
        }
    }

    fn binop(&mut self, binop: &ast::BinOp, lhs: &ast::Expression, rhs: &ast::Expression) -> Expression {
        use Expression::*;
        let operation = Builtin::from_str(binop.sigil()).ok();
        match operation {
            Some(Builtin::Assignment) => {
                let lval = match &lhs.kind {
                    ast::ExpressionKind::Value(qualified_name) => {
                        if let Some(Symbol { def_id, .. }) = self.symbol_table.lookup_symbol(&qualified_name.id) {
                            def_id.clone()
                        } else {
                            return ReductionError(format!("Couldn't look up name: {:?}", qualified_name));
                        }
                    },
                    _ => return ReductionError("Trying to assign to a non-name".to_string()),
                };

                Assign {
                    lval,
                    rval: Box::new(self.expression(rhs)),
                }
            },
            Some(op) => {
                Expression::Call {
                    f: Box::new(Expression::Callable(Function::Builtin(op))),
                    args: vec![self.expression(lhs), self.expression(rhs)],
                }
            }
            None => {
                //TODO handle a user-defined operation
                Unimplemented
            }
        }
    }

    fn value(&mut self, qualified_name: &ast::QualifiedName) -> Expression {
        use SymbolSpec::*;

        let ast::QualifiedName { id, components, .. } = qualified_name;

        let symbol = match self.symbol_table.lookup_symbol(&qualified_name.id) {
            Some(s) => s,
            None => return Expression::ReductionError(format!("No symbol found for name: {:?}", qualified_name))
        };
        let Symbol { def_id, spec, .. } = symbol;
        match spec {
            Func(_) => Expression::Lookup { id: def_id.clone(), kind: Lookup::Function },
            GlobalBinding => Expression::Lookup { id: def_id.clone(), kind: Lookup::GlobalVar },
            LocalVariable => Expression::Lookup { id: def_id.clone(), kind: Lookup::LocalVar },
            FunctionParam(_) => Expression::Lookup { id: def_id.clone(), kind: Lookup::Param },
            DataConstructor { index, arity, .. } => {
                Expression::Unimplemented
            },
            RecordConstructor { .. } => {
                Expression::ReductionError(format!("The symbol for value {:?} is unexpectdly a RecordConstructor", qualified_name))
            },
        }
    }
}

/// The reduced intermediate representation consists of a list of function definitions, and a block
/// of entrypoint statements.  In a repl or script context this can be an arbitrary list of
/// statements, in an executable context will likely just be a pointer to the main() function.
#[derive(Debug)]
pub struct ReducedIR {
    pub functions: HashMap<DefId, FunctionDefinition>,
    pub entrypoint: Vec<Statement>,
}

impl ReducedIR {
    fn debug(&self, symbol_table: &SymbolTable) {
        println!("Reduced IR:");
        println!("Functions:");
        println!("-----------");
        for (id, callable) in self.functions.iter() {
            let name = &symbol_table.lookup_symbol_by_def(id).unwrap().local_name;
            println!("{}({}) -> {:?}", id, name, callable);
        }
        println!("");
        println!("Entrypoint:");
        println!("-----------");
        for stmt in self.entrypoint.iter() {
            println!("{:?}", stmt);
        }
        println!("-----------");
    }
}

#[derive(Debug)]
pub enum Statement {
    Expression(Expression),
    Binding {
        id: DefId,
        constant: bool,
        expr: Expression
    },
}

#[derive(Debug)]
pub enum Expression {
    Literal(Literal),
    Tuple(Vec<Expression>),
    Lookup {
        id: DefId, //TODO eventually not everything that can be looked up will have a DefId
        kind: Lookup,
    },
    Assign {
        lval: DefId,
        rval: Box<Expression>,
    },
    Callable(Function),
    Call {
        f: Box<Expression>,
        args: Vec<Expression>
    },
    Unimplemented,
    ReductionError(String),
}

impl Expression {
    pub fn unit() -> Self {
        Expression::Tuple(vec![])
    }
}

#[derive(Debug)]
pub struct FunctionDefinition {
    pub body: Vec<Statement>
}

#[derive(Debug, Clone)]
pub enum Function {
    Builtin(Builtin),
    UserDefined(DefId)
}

#[derive(Debug)]
pub enum Lookup {
    LocalVar,
    GlobalVar,
    Function,
    Param,
}

#[derive(Debug, Clone)]
pub enum Literal {
    Nat(u64),
    Int(i64),
    Float(f64),
    Bool(bool),
    StringLit(Rc<String>),
}