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

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

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

mod types;
mod test;

pub use types::*;

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();
                    let def_id = symbol.def_id().unwrap();
                    entrypoint.push(Statement::Binding { id: def_id, 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
                },
                ast::Declaration::Binding { constant, expr, ..} => {
                    let symbol = self.symbol_table.lookup_symbol(item_id).unwrap();
                    let def_id = symbol.def_id().unwrap();
                    Some(Statement::Binding { id: def_id, constant: *constant, expr: self.expression(&expr) })
                },

                _ => 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().unwrap();
        let function_def = FunctionDefinition {
            body: self.function_internal_block(statements)
        };
        self.functions.insert(def_id, function_def);
    }

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

        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 } =>  Expression::Call {
                f: Box::new(self.expression(f)),
                args: arguments
                    .iter()
                    .map(|arg| self.invocation_argument(arg))
                    .collect(),
            },
            TupleLiteral(exprs) => Expression::Tuple(exprs.iter().map(|e| self.expression(e)).collect()),
            IfExpression { discriminator, body, } => self.reduce_if_expression(discriminator.as_ref().map(|x| x.as_ref()), body),
            Lambda { params, body, .. } =>  {
                Expression::Callable(Callable::Lambda {
                    arity: params.len() as u8,
                    body: self.function_internal_block(body),
                })
            },
            NamedStruct { .. } => Expression::ReductionError("NamedStruct not implemented".to_string()), //self.reduce_named_struct(name, fields),
            Index { .. } => Expression::ReductionError("Index expr not implemented".to_string()),
            WhileExpression { .. } => Expression::ReductionError("While expr not implemented".to_string()),
            ForExpression { .. } => Expression::ReductionError("For expr not implemented".to_string()),
            ListLiteral { .. } => Expression::ReductionError("ListLiteral expr not implemented".to_string()),
        }
    }

    fn reduce_if_expression(&mut self, discriminator: Option<&ast::Expression>, body: &ast::IfExpressionBody) -> Expression {
        use  ast::IfExpressionBody::*;

        let cond = Box::new(match discriminator {
            Some(expr) => self.expression(expr),
            None => return Expression::ReductionError("blank cond if-expr not supported".to_string()),
        });
        match body {
            SimpleConditional {
                then_case,
                else_case,
            } => {
                let then_clause = self.function_internal_block(then_case);
                let else_clause = match else_case.as_ref() {
                    None => vec![],
                    Some(stmts) => self.function_internal_block(stmts),
                };
                Expression::Conditional {
                    cond,
                    then_clause,
                    else_clause,
                }
            },
            SimplePatternMatch {
                pattern,
                then_case,
                else_case,
            } => {
                let alternatives = vec![
                    Alternative {
                        pattern: match pattern.reduce(self.symbol_table) {
                            Ok(p) => p,
                            Err(e) => return Expression::ReductionError(format!("Bad pattern: {:?}", e)),
                        },
                        item: self.function_internal_block(then_case),
                    },
                    Alternative {
                        pattern: Pattern::Ignored,
                        item: match else_case.as_ref() {
                            Some(else_case) => self.function_internal_block(else_case),
                            None => vec![],
                        },
                    },
                ];

                Expression::CaseMatch { cond, alternatives }
            },
            CondList(ref condition_arms) => {
                let mut alternatives = vec![];
                for arm in condition_arms {
                    match arm.condition {
                        ast::Condition::Expression(ref _expr) => return Expression::ReductionError("case-expression".to_string()),
                        ast::Condition::Pattern(ref pat) => {
                            let alt = Alternative {
                                pattern: match pat.reduce(self.symbol_table) {
                                    Ok(p) => p,
                                    Err(e) => return Expression::ReductionError(format!("Bad pattern: {:?}", e)),
                                },
                                item: self.function_internal_block(&arm.body),
                            };
                            alternatives.push(alt);
                        },
                        ast::Condition::TruncatedOp(_, _) =>  return Expression::ReductionError("case-expression-trunc-op".to_string()),
                        ast::Condition::Else => return Expression::ReductionError("case-expression-else".to_string()),
                    }
                }
                Expression::CaseMatch { cond, alternatives }
            }
        }
    }

    fn invocation_argument(&mut self, invoc: &ast::InvocationArgument) -> Expression {
        use crate::ast::InvocationArgument::*;
        match invoc {
            Positional(ex) => self.expression(ex),
            Keyword { .. } => Expression::ReductionError("Keyword arguments not supported".to_string()),
            Ignored  => Expression::ReductionError("Ignored arguments not supported".to_string()),
        }
    }

    fn function_internal_block(&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(Callable::Builtin(op))),
                    args: vec![self.expression(arg)],
                }
            }
            None => {
                //TODO need this for custom prefix ops
                Expression::ReductionError("User-defined prefix ops not supported".to_string())
            }
        }
    }

    fn binop(&mut self, binop: &ast::BinOp, lhs: &ast::Expression, rhs: &ast::Expression) -> Expression {
        use Expression::ReductionError;

        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) = self.symbol_table.lookup_symbol(&qualified_name.id) {
                            symbol.def_id().unwrap()
                        } else {
                            return ReductionError(format!("Couldn't look up name: {:?}", qualified_name));
                        }
                    },
                    _ => return ReductionError("Trying to assign to a non-name".to_string()),
                };

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

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

        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 def_id = symbol.def_id();

        match symbol.spec() {
            Func => Expression::Lookup(Lookup::Function(def_id.unwrap())),
            GlobalBinding => Expression::Lookup(Lookup::GlobalVar(def_id.unwrap())),
            LocalVariable => Expression::Lookup(Lookup::LocalVar(def_id.unwrap())),
            FunctionParam(n) => Expression::Lookup(Lookup::Param(n)),
            DataConstructor { index, arity, type_id } => Expression::Callable(Callable::DataConstructor {
                type_id: type_id.clone(),
                arity: arity as u32, //TODO fix up these modifiers
                tag: index as u32,
            }),
            RecordConstructor { .. } => {
                Expression::ReductionError(format!("The symbol for value {:?} is unexpectdly a RecordConstructor", qualified_name))
            },
        }
    }
}

impl ast::Pattern {
    fn reduce(&self, symbol_table: &SymbolTable) -> Result<Pattern, PatternError> {
        Ok(match self {
            ast::Pattern::Ignored => {
                Pattern::Ignored
            },
            ast::Pattern::TuplePattern(subpatterns) => {
                let pats: Vec<_> = subpatterns.iter().map(|pat| pat.reduce(symbol_table)).collect();
                let pats: Result<Vec<Pattern>, PatternError> = pats.into_iter().collect();
                let pats = pats?;
                Pattern::Tuple(pats)
            },
            ast::Pattern::Literal(lit) => Pattern::Literal(match lit {
                ast::PatternLiteral::NumPattern { neg, num } =>  match (neg, num) {
                    (false, ast::ExpressionKind::NatLiteral(n)) => Literal::Nat(*n),
                    (false, ast::ExpressionKind::FloatLiteral(f)) => Literal::Float(*f),
                    (true, ast::ExpressionKind::NatLiteral(n)) => Literal::Int(-(*n as i64)),
                    (true, ast::ExpressionKind::FloatLiteral(f)) => Literal::Float(-f),
                    (_, e) => return Err(format!("Internal error, unexpected pattern literal: {:?}", e).into())
                },
                ast::PatternLiteral::StringPattern(s) => Literal::StringLit(s.clone()),
                ast::PatternLiteral::BoolPattern(b) => Literal::Bool(*b),
            }),
            ast::Pattern::TupleStruct(_name, _subpatterns) => {
                unimplemented!()
            },
            ast::Pattern::Record(name, /*Vec<(Rc<String>, Pattern)>*/ _) => {
                unimplemented!()
            },
            ast::Pattern::VarOrName(_name) => {
                unimplemented!()
            },
        })
    }
}

/*
impl ast::Pattern {
    fn to_subpattern(&self, symbol_table: &SymbolTable) -> Subpattern {
        use ast::Pattern::*;
        use Expression::ReductionError;

        match self {
            Ignored => Subpattern {
                tag: None,
                subpatterns: vec![],
                guard: None,
            },
            Literal(lit) => lit.to_subpattern(symbol_table),
            /*
            TupleStruct(QualifiedName { components, id }, inner_patterns) => {
                match symbol_table.lookup_symbol(id) {
                    Some(symbol) => handle_symbol(Some(symbol), inner_patterns, symbol_table),
                    None => panic!("Symbol {:?} not found", components),
                }
            }
            */
            _ => Subpattern {
                tag: None,
                subpatterns: vec![],
                guard: None,
            }
        }
            /*
            TuplePattern(inner_patterns) => handle_symbol(None, inner_patterns, symbol_table),
            Record(_name, _pairs) => {
                unimplemented!()
            }
            VarOrName(QualifiedName { components, id }) => {
                // if symbol is Some, treat this as a symbol pattern. If it's None, treat it
                // as a variable.
                match symbol_table.lookup_symbol(id) {
                    Some(symbol) => handle_symbol(Some(symbol), &[], symbol_table),
                    None => {
                        println!("Components: {:?}", components);
                        let name = if components.len() == 1 {
                            components[0].clone()
                        } else {
                            panic!("check this line of code yo");
                        };
                        Subpattern {
                            tag: None,
                            subpatterns: vec![],
                            guard: None,
                            bound_vars: vec![Some(name)],
                        }
                    }
                }
            }
            */
    }
}
*/

/*
fn handle_symbol(
    symbol: Option<&Symbol>,
    inner_patterns: &[Pattern],
    symbol_table: &SymbolTable,
) -> Subpattern {
    use ast::Pattern::*;

    let tag = symbol.map(|symbol| match symbol.spec {
        SymbolSpec::DataConstructor { index, .. } => index,
        _ => {
            panic!("Symbol is not a data constructor - this should've been caught in type-checking")
        }
    });
    let bound_vars = inner_patterns
        .iter()
        .map(|p| match p {
            VarOrName(qualified_name) => {
                let symbol_exists = symbol_table.lookup_symbol(&qualified_name.id).is_some();
                if symbol_exists {
                    None
                } else {
                    let QualifiedName { components, .. } = qualified_name;
                    if components.len() == 1 {
                        Some(components[0].clone())
                    } else {
                        panic!("Bad variable name in pattern");
                    }
                }
            }
            _ => None,
        })
        .collect();

    let subpatterns = inner_patterns
        .iter()
        .map(|p| match p {
            Ignored => None,
            VarOrName(_) => None,
            Literal(other) => Some(other.to_subpattern(symbol_table)),
            tp @ TuplePattern(_) => Some(tp.to_subpattern(symbol_table)),
            ts @ TupleStruct(_, _) => Some(ts.to_subpattern(symbol_table)),
            Record(..) => unimplemented!(),
        })
        .collect();

    let guard = None;
    /*
    let guard_equality_exprs: Vec<Expr> = subpatterns.iter().map(|p| match p {
    Literal(lit) => match lit {
    _ => unimplemented!()
    },
    _ => unimplemented!()
    }).collect();
    */

    Subpattern {
        tag,
        subpatterns,
        guard,
        bound_vars,
    }
}
*/

/*
impl ast::PatternLiteral {
    fn to_subpattern(&self, _symbol_table: &SymbolTable) -> Subpattern {
        use ast::PatternLiteral::*;

        match self {
            NumPattern { neg, num } => {
                let comparison = Expression::Literal(match (neg, num) {
                    (false, ast::ExpressionKind::NatLiteral(n)) => Literal::Nat(*n),
                    (false, ast::ExpressionKind::FloatLiteral(f)) => Literal::Float(*f),
                    (true, ast::ExpressionKind::NatLiteral(n)) => Literal::Int(-(*n as i64)),
                    (true, ast::ExpressionKind::FloatLiteral(f)) => Literal::Float(-f),
                    _ => panic!("This should never happen"),
                });
                unimplemented!()
                /*
                let guard = Some(Expr::Call {
                    f: Box::new(Expr::Func(Func::BuiltIn(Builtin::Equality))),
                    args: vec![comparison, Expr::ConditionalTargetSigilValue],
                });
                Subpattern {
                    tag: None,
                    subpatterns: vec![],
                    guard,
                }
                */
            }
            _ => unimplemented!()
            /*
            StringPattern(s) => {
                let guard = Some(Expr::Call {
                    f: Box::new(Expr::Func(Func::BuiltIn(Builtin::Equality))),
                    args: vec![
                        Expr::Lit(Lit::StringLit(s.clone())),
                        Expr::ConditionalTargetSigilValue,
                    ],
                });

                Subpattern {
                    tag: None,
                    subpatterns: vec![],
                    guard,
                    bound_vars: vec![],
                }
            }
            BoolPattern(b) => {
                let guard = Some(if *b {
                    Expr::ConditionalTargetSigilValue
                } else {
                    Expr::Call {
                        f: Box::new(Expr::Func(Func::BuiltIn(Builtin::BooleanNot))),
                        args: vec![Expr::ConditionalTargetSigilValue],
                    }
                });
                Subpattern {
                    tag: None,
                    subpatterns: vec![],
                    guard,
                    bound_vars: vec![],
                }
            }
            */
        }
    }
}
*/