oden/fine/src/semantics.rs

use crate::{
    parser::{Child, SyntaxTree, Tree, TreeKind, TreeRef},
    tokens::{Lines, Token, TokenKind},
};
use std::{cell::RefCell, collections::HashMap, fmt, rc::Rc};

// TODO: Unused variables?

// TODO: An error should have:
//
//   - a start
//   - an end
//   - a focus
//   - descriptive messages
//
// that will have to wait for now
#[derive(Clone, PartialEq, Eq)]
pub struct Error {
    pub start: (usize, usize),
    pub end: (usize, usize),
    pub message: String,
}

impl Error {
    pub fn new<T>(line: usize, column: usize, message: T) -> Self
    where
        T: ToString,
    {
        Error {
            start: (line, column),
            end: (line, column),
            message: message.to_string(),
        }
    }

    pub fn new_spanned<T>(start: (usize, usize), end: (usize, usize), message: T) -> Self
    where
        T: ToString,
    {
        Error {
            start,
            end,
            message: message.to_string(),
        }
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{self}")
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{}: {}", self.start.0, self.start.1, self.message)
    }
}

#[derive(Clone)]
pub enum Type {
    // Signals a type error. If you receive this then you know that an error
    // has already been reported; if you produce this be sure to also note
    // the error in the errors collection.
    Error,

    // Signals that the expression has a control-flow side-effect and that no
    // value will ever result from this expression. Usually this means
    // everything's fine.
    Unreachable,

    // This is until generics are working
    MagicPrintGarbage,

    Nothing,
    // TODO: Numeric literals should be implicitly convertable, unlike other
    //       types. Maybe just "numeric literal" type?
    F64,
    String,
    Bool,

    Function(Vec<Box<Type>>, Box<Type>),
}

impl Type {
    pub fn is_error(&self) -> bool {
        match self {
            Type::Error => true,
            _ => false,
        }
    }

    pub fn compatible_with(&self, other: &Type) -> bool {
        // TODO: This is wrong; we because of numeric literals etc.
        match (self, other) {
            (Type::F64, Type::F64) => true,
            (Type::String, Type::String) => true,
            (Type::Bool, Type::Bool) => true,
            (Type::Unreachable, Type::Unreachable) => true,
            (Type::Nothing, Type::Nothing) => true,

            // Avoid introducing more errors
            (Type::Error, _) => true,
            (_, Type::Error) => true,

            (_, _) => false,
        }
    }
}

impl fmt::Debug for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{self}")
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use Type::*;
        match self {
            Error => write!(f, "<< INTERNAL ERROR >>"),
            Unreachable => write!(f, "<< UNREACHABLE >>"),
            Nothing => write!(f, "()"),
            F64 => write!(f, "f64"),
            String => write!(f, "string"),
            Bool => write!(f, "bool"),
            MagicPrintGarbage => write!(f, "MagicPrintGarbage"),
            Function(args, ret) => {
                write!(f, "(")?;
                let mut first = true;
                for arg in args.iter() {
                    if !first {
                        write!(f, ", ")?;
                    }
                    write!(f, "{arg}")?;
                    first = false;
                }
                write!(f, ") -> {ret}")
            }
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub enum Location {
    Argument,
    Local,
    Module,
}

// TODO: Is `usize` what we want? Do we want e.g. dyn trait for invoke?
#[derive(Clone, Debug)]
pub struct ExternalFunctionId(usize);

impl ExternalFunctionId {
    pub fn id(&self) -> usize {
        self.0
    }
}

pub enum Declaration {
    Variable {
        declaration_type: Type,
        location: Location,
        index: usize,
    },
    Function {
        declaration_type: Type,
        declaration: TreeRef, //?
    },
    ExternFunction {
        declaration_type: Type,
        id: ExternalFunctionId,
    },
}

pub struct Environment {
    pub parent: Option<EnvironmentRef>,
    pub location: Location,
    pub next_index: usize,
    pub declarations: HashMap<Box<str>, Declaration>,
}

impl Environment {
    pub fn new(parent: Option<EnvironmentRef>, location: Location) -> Self {
        let parent_location = parent
            .as_ref()
            .map(|p| p.location)
            .unwrap_or(Location::Module);
        let base = parent.as_ref().map(|p| p.next_index).unwrap_or(0);
        let next_index = match (parent_location, location) {
            (_, Location::Argument) => 0,

            (Location::Local, Location::Local) => base,
            (_, Location::Local) => 0,

            (Location::Module, Location::Module) => base,
            (_, Location::Module) => panic!("What?"),
        };

        Environment {
            parent,
            location,
            next_index,
            declarations: HashMap::new(),
        }
    }

    pub fn insert(&mut self, token: &Token, t: Type) {
        self.declarations.insert(
            token.as_str().into(),
            Declaration::Variable {
                declaration_type: t,
                location: self.location,
                index: self.next_index,
            },
        );
        self.next_index += 1;
    }

    pub fn bind(&self, token: &Token) -> Option<&Declaration> {
        if let Some(decl) = self.declarations.get(token.as_str()) {
            return Some(decl);
        }

        let mut current = &self.parent;
        while let Some(env) = current {
            if let Some(decl) = env.declarations.get(token.as_str()) {
                return Some(decl);
            }
            current = &env.parent;
        }

        None
    }
}

#[derive(Clone)]
pub struct EnvironmentRef(Rc<Environment>);

impl EnvironmentRef {
    pub fn new(environment: Environment) -> Self {
        EnvironmentRef(Rc::new(environment))
    }
}

impl std::ops::Deref for EnvironmentRef {
    type Target = Environment;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

fn set_logical_parents(
    parents: &mut Vec<Option<TreeRef>>,
    syntax_tree: &SyntaxTree,
    t: TreeRef,
    parent: Option<TreeRef>,
) {
    parents[t.index()] = parent.clone();

    let tree = &syntax_tree[t];
    // eprintln!("SET PARENT {parent:?} => CHILD {tree:?} ({t:?})");
    match tree.kind {
        TreeKind::Block | TreeKind::File => {
            // In a block (or at the top level), each child actually points
            // to the previous child as the logical parent, so that variable
            // declarations that occur as part of statements in the block are
            // available to statements later in the block.
            let mut parent = Some(t);
            for child in &tree.children {
                match child {
                    Child::Token(_) => (),
                    Child::Tree(ct) => {
                        set_logical_parents(parents, syntax_tree, *ct, parent);
                        parent = Some(*ct);
                    }
                }
            }
        }
        TreeKind::LetStatement => {
            // In a let statement, the logical parent of the children is
            // actually the logical parent of the let statement, so that the
            // variable doesn't have itself in scope. :P
            for child in &tree.children {
                match child {
                    Child::Token(_) => (),
                    Child::Tree(ct) => set_logical_parents(parents, syntax_tree, *ct, parent),
                }
            }
        }
        TreeKind::FunctionDecl => {
            // In a function declaration, the logical parent of the body is
            // the parameter list.
            let param_list = tree.child_of_kind(syntax_tree, TreeKind::ParamList);
            let body = tree.child_of_kind(syntax_tree, TreeKind::Block);
            for child in &tree.children {
                match child {
                    Child::Token(_) => (),
                    Child::Tree(ct) => {
                        if Some(*ct) == body {
                            set_logical_parents(parents, syntax_tree, *ct, param_list);
                        } else {
                            set_logical_parents(parents, syntax_tree, *ct, Some(t));
                        }
                    }
                }
            }
        }
        _ => {
            // By default, the parent for each child is current tree.
            for child in &tree.children {
                match child {
                    Child::Token(_) => (),
                    Child::Tree(ct) => set_logical_parents(parents, syntax_tree, *ct, Some(t)),
                }
            }
        }
    }
}

#[derive(Clone, Copy, Debug)]
enum Incremental<T> {
    None,
    InProgress,
    Complete(T),
}

pub struct Semantics<'a> {
    // TODO: Do I really want my own copy here? Should we standardize on Arc
    //       or Rc or some other nice sharing mechanism?
    syntax_tree: &'a SyntaxTree<'a>,
    lines: &'a Lines,

    // Instead of physical parents, this is the set of *logical* parents.
    // This is what is used for binding.
    logical_parents: Vec<Option<TreeRef>>,

    // TODO: State should be externalized instead of this refcell nonsense.
    errors: RefCell<Vec<Error>>,
    types: RefCell<Vec<Incremental<Type>>>,
    environments: RefCell<Vec<Incremental<EnvironmentRef>>>,
    root_environment: EnvironmentRef,
}

impl<'a> Semantics<'a> {
    pub fn new(tree: &'a SyntaxTree<'a>, lines: &'a Lines) -> Self {
        let mut logical_parents = vec![None; tree.len()];
        if let Some(root) = tree.root() {
            set_logical_parents(&mut logical_parents, tree, root, None);
        }

        let mut root_environment = Environment::new(None, Location::Module);
        root_environment.declarations.insert(
            "print".into(),
            Declaration::ExternFunction {
                declaration_type: Type::MagicPrintGarbage,
                id: ExternalFunctionId(0),
            },
        );

        let mut semantics = Semantics {
            syntax_tree: tree,
            lines,
            logical_parents,
            errors: RefCell::new(vec![]),
            types: RefCell::new(vec![Incremental::None; tree.len()]),
            environments: RefCell::new(vec![Incremental::None; tree.len()]),
            root_environment: EnvironmentRef::new(root_environment),
        };

        // NOTE: We ensure all the known errors are reported before we move
        //       on to answering any other questions. We're going to work as
        //       hard as we can from a partial tree.
        if let Some(tr) = semantics.syntax_tree.root() {
            semantics.gather_errors(tr);
        }

        semantics
    }

    pub fn tree(&self) -> &SyntaxTree<'a> {
        &self.syntax_tree
    }

    pub fn snapshot_errors(&self) -> Vec<Error> {
        (*self.errors.borrow()).clone()
    }

    pub fn logical_parent(&self, tr: TreeRef) -> Option<TreeRef> {
        if tr.index() < self.logical_parents.len() {
            self.logical_parents[tr.index()]
        } else {
            None
        }
    }

    fn report_error<T>(&self, position: usize, error: T)
    where
        T: ToString,
    {
        let (line, col) = self.lines.position(position);
        self.errors
            .borrow_mut()
            .push(Error::new(line, col, error.to_string()));
    }

    fn report_error_span<T>(&self, start: usize, end: usize, error: T)
    where
        T: ToString,
    {
        let start = self.lines.position(start);
        let end = self.lines.position(end);
        self.errors
            .borrow_mut()
            .push(Error::new_spanned(start, end, error.to_string()));
    }

    fn report_error_tree<T>(&self, tree: &Tree<'a>, error: T)
    where
        T: ToString,
    {
        self.report_error_span(tree.start_pos, tree.end_pos, error)
    }

    fn report_error_tree_ref<T>(&self, tree: TreeRef, error: T)
    where
        T: ToString,
    {
        let tree = &self.syntax_tree[tree];
        self.report_error_span(tree.start_pos, tree.end_pos, error)
    }

    fn gather_errors(&mut self, tree: TreeRef) {
        let mut stack = vec![tree];
        while let Some(tr) = stack.pop() {
            let tree = &self.syntax_tree[tr];
            for child in &tree.children {
                match child {
                    Child::Token(t) => {
                        if t.kind == TokenKind::Error {
                            self.report_error(t.start, t.as_str());
                        }
                    }
                    Child::Tree(t) => stack.push(*t),
                }
            }
        }
    }

    pub fn environment_of(&self, t: TreeRef) -> EnvironmentRef {
        {
            // I want to make sure that this borrow is dropped after this block.
            let mut borrow = self.environments.borrow_mut();
            let state = &mut borrow[t.index()];
            match state {
                Incremental::None => (),
                Incremental::Complete(e) => return e.clone(),
                Incremental::InProgress => {
                    // NOTE: Set the state so the ICE doesn't loop on itself.
                    *state = Incremental::Complete(self.root_environment.clone());
                    drop(borrow);

                    //eprintln!("environment_of circular => {t:?}");
                    self.internal_compiler_error(Some(t), "circular environment dependency");
                }
            }
            *state = Incremental::InProgress;
        }

        let tree = &self.syntax_tree[t];
        //eprintln!(">>> environment_of => {tree:?}");

        let parent = match self.logical_parents[t.index()] {
            Some(t) => self.environment_of(t),
            None => self.root_environment.clone(),
        };

        let result = match tree.kind {
            TreeKind::LetStatement => self.environment_of_let(parent, tree),
            TreeKind::ParamList => self.environment_of_paramlist(parent, tree),
            TreeKind::File => self.environment_of_file(parent, tree),
            TreeKind::Block => self.environment_of_block(parent, tree),

            // TODO: Blocks should introduce a local environment if required.
            //       Test with a type error in a block statement and a
            //       binding outside. You will need a new assertion type and
            //       possibly a compile/run to ensure it works.
            //
            //  let x = 7;
            //  {
            //    let x = 23;
            //  }
            //  print(x); // 7
            //
            //  {
            //     let y = 12; // check: `y` is local not global!
            //  }
            //  print(y); // error, cannot find 'y'

            // TODO: MORE Things that introduce an environment!
            _ => parent,
        };

        self.environments.borrow_mut()[t.index()] = Incremental::Complete(result.clone());
        //eprintln!("<<< environment_of => {tree:?}");
        result
    }

    fn environment_of_block(&self, parent: EnvironmentRef, tree: &Tree) -> EnvironmentRef {
        let mut environment = Environment::new(Some(parent), Location::Local);
        for child in tree.children.iter() {
            match child {
                Child::Tree(t) => {
                    let ct = &self.syntax_tree[*t];
                    if ct.kind == TreeKind::FunctionDecl {
                        // TODO: Should I have accessors for function decls?
                        let Some(name) = ct.nth_token(1) else {
                            continue;
                        };

                        environment.declarations.insert(
                            name.as_str().into(),
                            Declaration::Function {
                                declaration_type: self.type_of(*t),
                                declaration: *t,
                            },
                        );
                    }
                }
                _ => {}
            }
        }

        EnvironmentRef::new(environment)
    }

    fn environment_of_file(&self, parent: EnvironmentRef, tree: &Tree) -> EnvironmentRef {
        let mut environment = Environment::new(Some(parent), Location::Module);
        for child in tree.children.iter() {
            match child {
                Child::Tree(t) => {
                    let ct = &self.syntax_tree[*t];
                    if ct.kind == TreeKind::FunctionDecl {
                        // TODO: Should I have accessors for function decls?
                        let Some(name) = ct.nth_token(1) else {
                            continue;
                        };

                        environment.declarations.insert(
                            name.as_str().into(),
                            Declaration::Function {
                                declaration_type: self.type_of(*t),
                                declaration: *t,
                            },
                        );
                    }
                }
                _ => {}
            }
        }

        EnvironmentRef::new(environment)
    }

    fn environment_of_let(&self, parent: EnvironmentRef, tree: &Tree) -> EnvironmentRef {
        let Some(name) = tree.nth_token(1) else {
            return parent; // Error is already reported?
        };

        let declaration_type = match tree.nth_tree(3) {
            Some(expr) => self.type_of(expr),

            // The syntax error should already have been reported, so we'll
            // stick with error type here. (But bind the name, because we see
            // it!)
            None => Type::Error,
        };

        let location = match parent.location {
            Location::Local => Location::Local,
            Location::Module => Location::Module,
            Location::Argument => Location::Local,
        };

        let mut environment = Environment::new(Some(parent), location);
        environment.insert(name, declaration_type);

        EnvironmentRef::new(environment)
    }

    fn environment_of_paramlist(&self, parent: EnvironmentRef, tree: &Tree) -> EnvironmentRef {
        assert!(tree.kind == TreeKind::ParamList);

        let mut environment = Environment::new(Some(parent), Location::Argument);
        for child in tree.children.iter() {
            let Child::Tree(ct) = child else {
                continue;
            };
            let param = &self.syntax_tree[*ct];
            if param.kind != TreeKind::Parameter {
                continue;
            }

            let Some(param_name) = param.nth_token(0) else {
                continue;
            };

            let declaration_type = if let Some(type_expression) = param.nth_tree(2) {
                self.type_of(type_expression)
            } else {
                Type::Error
            };

            environment.insert(param_name, declaration_type);
        }

        EnvironmentRef::new(environment)
    }

    pub fn type_of(&self, t: TreeRef) -> Type {
        {
            let state = &mut self.types.borrow_mut()[t.index()];
            match state {
                Incremental::None => (),
                Incremental::Complete(existing) => return existing.clone(),
                Incremental::InProgress => {
                    // eprintln!("type_of circular => {t:?}");
                    self.report_error_tree_ref(t, "The type of this expression depends on itself");
                    *state = Incremental::Complete(Type::Error);
                    return Type::Error;
                }
            }
            *state = Incremental::InProgress;
        }

        let tree = &self.syntax_tree[t];
        // eprintln!("type_of => {tree:?}");

        let result = match tree.kind {
            TreeKind::Error => Some(Type::Error),
            TreeKind::UnaryExpression => self.type_of_unary(tree),
            TreeKind::BinaryExpression => self.type_of_binary(tree),
            TreeKind::TypeExpression => self.type_of_type_expr(tree),
            TreeKind::Block => self.type_of_block(tree),
            TreeKind::LiteralExpression => self.type_of_literal(tree),
            TreeKind::GroupingExpression => self.type_of_grouping(tree),
            TreeKind::ConditionalExpression => self.type_of_conditional(tree),
            TreeKind::CallExpression => self.type_of_call(tree),
            TreeKind::Argument => self.type_of_argument(tree),

            TreeKind::LetStatement => Some(Type::Nothing),
            TreeKind::ReturnStatement => Some(Type::Unreachable),
            TreeKind::ExpressionStatement => self.type_of_expression_statement(tree),
            TreeKind::IfStatement => self.type_of_if_statement(tree),
            TreeKind::Identifier => self.type_of_identifier(t, tree),

            TreeKind::FunctionDecl => self.type_of_function_decl(tree),
            TreeKind::ReturnType => self.type_of_return_type(tree),

            _ => self.internal_compiler_error(Some(t), "asking for a nonsense type"),
        };

        // NOTE: These return `None` if they encounter some problem.
        let result = result.unwrap_or(Type::Error);

        if result.is_error() {
            eprintln!("OH NO AN ERROR AT {}: {:?}", t.index(), tree);
        }

        self.types.borrow_mut()[t.index()] = Incremental::Complete(result.clone());
        result
    }

    fn type_of_unary(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::UnaryExpression);

        let op = tree.nth_token(0)?;
        let expr = tree.nth_tree(1)?;

        let argument_type = self.type_of(expr);
        match (op.kind, argument_type) {
            (TokenKind::Plus, Type::F64) => Some(Type::F64),
            (TokenKind::Minus, Type::F64) => Some(Type::F64),
            (TokenKind::Bang, Type::Bool) => Some(Type::Bool),

            // This is dumb and should be punished, probably.
            (_, Type::Unreachable) => {
                self.report_error(
                    op.start,
                    "cannot apply a unary operator to something that doesn't yield a value",
                );
                Some(Type::Error)
            }

            // Propagate existing errors without additional complaint.
            (_, Type::Error) => Some(Type::Error),

            (_, arg_type) => {
                self.report_error(
                    op.start,
                    format!(
                        "cannot apply unary operator '{}' to value of type {}",
                        op.as_str(),
                        arg_type
                    ),
                );
                Some(Type::Error)
            }
        }
    }

    fn type_of_binary(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::BinaryExpression);
        let lhs = self.type_of(tree.nth_tree(0)?);
        let op = tree.nth_token(1)?;
        let rhs = self.type_of(tree.nth_tree(2)?);

        match (op.kind, lhs, rhs) {
            (
                TokenKind::Plus | TokenKind::Minus | TokenKind::Star | TokenKind::Slash,
                Type::F64,
                Type::F64,
            ) => Some(Type::F64),

            (TokenKind::Plus, Type::String, Type::String) => Some(Type::String),

            (TokenKind::And | TokenKind::Or, Type::Bool, Type::Bool) => Some(Type::Bool),

            // This is dumb and should be punished, probably.
            (_, _, Type::Unreachable) => {
                self.report_error(
                    op.start,
                    format!("cannot apply '{op}' to an argument that doesn't yield a value (on the right)"),
                );
                Some(Type::Error)
            }
            (_, Type::Unreachable, _) => {
                self.report_error(
                    op.start,
                    format!("cannot apply '{op}' to an argument that doesn't yield a value (on the left)"),
                );
                Some(Type::Error)
            }

            // Propagate existing errors without additional complaint.
            (_, Type::Error, _) => Some(Type::Error),
            (_, _, Type::Error) => Some(Type::Error),

            // Missed the whole table, it must be an error.
            (_, left_type, right_type) => {
                self.report_error(
                    op.start,
                    format!("cannot apply binary operator '{op}' to expressions of type '{left_type}' (on the left) and '{right_type}' (on the right)"),
                );
                Some(Type::Error)
            }
        }
    }

    fn type_of_type_expr(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::TypeExpression);

        // TODO: This will *clearly* need to get better.
        let token = tree.nth_token(0)?;
        match token.as_str() {
            "f64" => Some(Type::F64),
            "string" => Some(Type::String),
            "bool" => Some(Type::Bool),
            "()" => Some(Type::Nothing),
            _ => {
                self.report_error_tree(tree, "Unrecognized type");
                Some(Type::Error)
            }
        }
    }

    fn type_of_block(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::Block);

        if tree.children.len() < 2 {
            return None;
        }

        if tree.children.len() == 2 {
            // Empty blocks generate Nothing.
            return Some(Type::Nothing);
        }

        // The type of the block is the type of the last expression.
        // (But the last child is the closing brace probably?)
        let last_is_brace = tree.nth_token(tree.children.len() - 1).is_some();
        let last_index = tree.children.len() - if last_is_brace { 2 } else { 1 };

        let mut is_unreachable = false;
        for i in 1..last_index {
            // TODO: if `is_unreachable` here then we actually have
            //       unreachable code here! We should warn about it
            //       I guess.
            is_unreachable =
                matches!(self.type_of(tree.nth_tree(i)?), Type::Unreachable) || is_unreachable;
        }

        // NOTE: If for some reason the last statement is unsuitable for a
        //       type then we consider the type of the block to be Nothing.
        //       (And explicitly not Error, which is what returning None
        //       would yield.)
        let last_type = self.type_of(tree.nth_tree(last_index)?);

        // If anything in this block generated an "Unreachable" then the
        // whole type of the block is "unreachable" no matter what.
        Some(if is_unreachable {
            Type::Unreachable
        } else {
            last_type
        })
    }

    fn type_of_literal(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::LiteralExpression);

        let tok = tree.nth_token(0)?;
        let pig = match tok.kind {
            TokenKind::Number => Type::F64,
            TokenKind::String => Type::String,
            TokenKind::True | TokenKind::False => Type::Bool,
            _ => panic!("the token {tok} doesn't have a type!"),
        };
        Some(pig)
    }

    fn type_of_grouping(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::GroupingExpression);

        tree.nth_tree(1).map(|t| self.type_of(t))
    }

    fn type_of_conditional(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::ConditionalExpression);

        let cond_tree = tree.nth_tree(1)?;
        let cond_type = self.type_of(cond_tree);
        let then_type = self.type_of(tree.nth_tree(2)?);

        let has_else = tree
            .nth_token(3)
            .map(|t| t.kind == TokenKind::Else)
            .unwrap_or(false);
        let else_type = if has_else {
            Some(self.type_of(tree.nth_tree(4)?))
        } else {
            None
        };

        if !cond_type.compatible_with(&Type::Bool) {
            if !cond_type.is_error() {
                self.report_error_tree_ref(cond_tree, "conditions must yield a boolean");
            }
            Some(Type::Error)
        } else {
            match (then_type, else_type) {
                (Type::Error, _) => Some(Type::Error),
                (_, Some(Type::Error)) => Some(Type::Error),

                (Type::Unreachable, None) => Some(Type::Unreachable),
                (Type::Unreachable, Some(t)) => Some(t),
                (t, Some(Type::Unreachable)) => Some(t),

                (then_type, else_type) => {
                    let else_type = else_type.unwrap_or(Type::Nothing);
                    if !then_type.compatible_with(&else_type) {
                        self.report_error_tree(
                            tree,
                            format!("the type of the `then` branch ({then_type}) must match the type of the `else` branch ({else_type})"),
                        );
                        Some(Type::Error)
                    } else {
                        Some(then_type)
                    }
                }
            }
        }
    }

    fn type_of_call(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::CallExpression);

        let f_ref = tree.nth_tree(0)?;
        let f = self.type_of(f_ref);

        let arg_list = &self.syntax_tree[tree.nth_tree(1)?];
        let arg_types: Vec<_> = arg_list
            .children
            .iter()
            .filter_map(|c| match c {
                Child::Tree(t) => Some((*t, self.type_of(*t))),
                _ => None,
            })
            .collect();

        if f.is_error() || arg_types.iter().any(|(_, t)| t.is_error()) {
            return Some(Type::Error);
        }

        match f {
            Type::Function(params, ret) => {
                let mut any_errors = false;
                if params.len() != arg_types.len() {
                    // TODO: Augment with function name if known
                    self.report_error_tree(tree, format!("expected {} parameters", params.len()));
                    any_errors = true;
                }

                for (i, ((t, a), p)) in arg_types.iter().zip(params.iter()).enumerate() {
                    if !a.compatible_with(p) {
                        self.report_error_tree_ref(
                            *t,
                            format!(
                                "parameter {i} has an incompatible type: expected {} but got {}",
                                p, a
                            ),
                        );
                        any_errors = true;
                    }
                }

                if any_errors {
                    return Some(Type::Error);
                }

                Some(*ret.clone())
            }
            Type::MagicPrintGarbage => {
                if arg_types.len() > 1 {
                    self.report_error_tree(tree, "print takes a single argument");
                    Some(Type::Error)
                } else if arg_types.len() == 0 {
                    Some(Type::Nothing)
                } else {
                    let mut arg_types = arg_types;
                    let (_, t) = arg_types.pop().unwrap();
                    Some(t)
                }
            }
            _ => {
                self.report_error_tree_ref(f_ref, format!("expected a function type, got: {f}"));
                Some(Type::Error)
            }
        }
    }

    fn type_of_argument(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::Argument);

        let result = self.type_of(tree.nth_tree(0)?);
        Some(result)
    }

    fn type_of_expression_statement(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::ExpressionStatement);
        let last_is_semicolon = tree
            .nth_token(tree.children.len() - 1)
            .map(|t| t.kind == TokenKind::Semicolon)
            .unwrap_or(false);

        let expression_type = self.type_of(tree.nth_tree(0)?);
        Some(match expression_type {
            Type::Unreachable => Type::Unreachable,
            _ => {
                // A semicolon at the end of an expression statement discards
                // the value, leaving us with nothing. (Even if the
                // expression otherwise generated a type error!)
                if last_is_semicolon {
                    Type::Nothing
                } else {
                    expression_type
                }
            }
        })
    }

    fn type_of_identifier(&self, t: TreeRef, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::Identifier);

        let id = tree.nth_token(0)?;
        let environment = self.environment_of(t);
        if let Some(declaration) = environment.bind(id) {
            return Some(match declaration {
                Declaration::Variable {
                    declaration_type, ..
                } => declaration_type.clone(),
                Declaration::Function {
                    declaration_type, ..
                } => declaration_type.clone(),
                Declaration::ExternFunction {
                    declaration_type, ..
                } => declaration_type.clone(),
            });
        }

        self.report_error_tree(tree, format!("cannot find value {id} here"));
        Some(Type::Error)
    }

    fn type_of_function_decl(&self, tree: &Tree) -> Option<Type> {
        let param_list = tree.child_tree_of_kind(self.syntax_tree, TreeKind::ParamList)?;
        let mut parameter_types = Vec::new();
        for p in param_list.child_trees() {
            let param = &self.syntax_tree[p];

            // TODO: Missing type expression means it's a generic function.
            let parameter_type = param.child_of_kind(self.syntax_tree, TreeKind::TypeExpression)?;
            parameter_types.push(Box::new(self.type_of(parameter_type)));
        }

        let return_type = match tree.child_tree_of_kind(self.syntax_tree, TreeKind::ReturnType) {
            Some(t) => self.type_of(t.child_of_kind(self.syntax_tree, TreeKind::TypeExpression)?),
            None => Type::Nothing,
        };
        let return_type = Box::new(return_type);
        Some(Type::Function(parameter_types, return_type))
    }

    fn type_of_return_type(&self, tree: &Tree) -> Option<Type> {
        assert_eq!(tree.kind, TreeKind::ReturnType);
        Some(self.type_of(tree.nth_tree(1)?)) // type expression
    }

    fn type_of_if_statement(&self, tree: &Tree) -> Option<Type> {
        Some(self.type_of(tree.nth_tree(0)?))
    }

    pub fn dump_compiler_state(&self, tr: Option<TreeRef>) {
        eprintln!("Parsed the tree as:");
        eprintln!("\n{}", self.syntax_tree.dump(true));

        {
            let errors = self.errors.borrow();
            if errors.len() == 0 {
                eprintln!("There were no errors reported during checking.\n");
            } else {
                eprintln!(
                    "{} error{} reported during checking:",
                    errors.len(),
                    if errors.len() == 1 { "" } else { "s" }
                );
                for error in errors.iter() {
                    eprintln!("  Error: {error}");
                }
                eprintln!();
            }
        }

        if let Some(tr) = tr {
            eprintln!("This is about the tree: {:?}", &self.syntax_tree[tr]);
            eprintln!("The logical parent chain of the tree was:\n");
            let mut current = Some(tr);
            while let Some(c) = current {
                let t = &self.syntax_tree[c];
                eprintln!("  {:?} [{}-{})", t.kind, t.start_pos, t.end_pos);
                current = self.logical_parents[c.index()];
            }
            eprintln!("\nThe environment of the tree was:");
            let mut environment = Some(self.environment_of(tr));
            while let Some(env) = environment {
                for (k, v) in env.declarations.iter() {
                    eprint!("  {k}: ");
                    match v {
                        Declaration::Variable {
                            declaration_type,
                            location,
                            index,
                        } => {
                            eprintln!("{declaration_type:?} (variable {location:?} {index})");
                        }
                        Declaration::Function {
                            declaration_type,
                            declaration,
                        } => {
                            eprintln!("{declaration_type:?} (function {declaration:?})");
                        }
                        Declaration::ExternFunction {
                            declaration_type,
                            id,
                        } => {
                            eprintln!("{declaration_type:?} (extern {id:?})");
                        }
                    };
                }
                environment = env.parent.clone();
            }
            eprintln!();
        }
    }

    pub fn internal_compiler_error(&self, tr: Option<TreeRef>, message: &str) -> ! {
        eprintln!("Internal compiler error: {message}!");
        self.dump_compiler_state(tr);
        panic!("INTERNAL COMPILER ERROR: {message}")
    }
}

pub fn check(s: &Semantics) {
    for t in s.syntax_tree.trees() {
        let tree = &s.syntax_tree[t];
        match tree.kind {
            TreeKind::Error => {} // already reported
            TreeKind::File => {}
            TreeKind::FunctionDecl => check_function_decl(s, t, tree),
            TreeKind::ParamList => {}
            TreeKind::Parameter => {}
            TreeKind::TypeExpression => {
                let _ = s.type_of_type_expr(tree);
            }
            TreeKind::Block => {
                let _ = s.type_of_block(tree);
            }
            TreeKind::LetStatement => {
                let _ = s.environment_of(t);
            }
            TreeKind::ReturnStatement => {}
            TreeKind::ExpressionStatement
            | TreeKind::LiteralExpression
            | TreeKind::GroupingExpression
            | TreeKind::UnaryExpression
            | TreeKind::ConditionalExpression
            | TreeKind::CallExpression
            | TreeKind::BinaryExpression => {
                let _ = s.type_of(t);
            }
            TreeKind::ArgumentList => {}
            TreeKind::Argument => {
                let _ = s.type_of(t);
            }
            TreeKind::IfStatement => {}
            TreeKind::Identifier => {
                let _ = s.type_of(t);
            }
            TreeKind::ReturnType => {}
        }
    }
}

fn check_function_decl(s: &Semantics, t: TreeRef, tree: &Tree) {
    assert_eq!(tree.kind, TreeKind::FunctionDecl);
    let _ = s.environment_of(t);

    let return_type_tree = tree.child_of_kind(s.syntax_tree, TreeKind::ReturnType);
    let return_type = return_type_tree
        .map(|t| s.type_of(t))
        .unwrap_or(Type::Nothing);

    if let Some(body) = tree.child_of_kind(s.syntax_tree, TreeKind::Block) {
        let body_type = s.type_of(body);
        if !body_type.compatible_with(&return_type) {
            // Just work very hard to get an appropriate error span.
            let (start, end) = return_type_tree
                .map(|t| {
                    let rtt = &s.syntax_tree[t];
                    (rtt.start_pos, rtt.end_pos)
                })
                .unwrap_or_else(|| {
                    let start = tree.start_pos;
                    let end_tok = tree
                        .nth_token(1)
                        .unwrap_or_else(|| tree.nth_token(0).unwrap());
                    let end_pos = end_tok.start + end_tok.as_str().len();
                    (start, end_pos)
                });

            s.report_error_span(start, end, format!("the body of this function yields a value of type `{body_type}`, but callers expect this function to produce a `{return_type}`"));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::parse;

    #[test]
    #[should_panic(expected = "INTERNAL COMPILER ERROR: oh no")]
    pub fn ice() {
        let (tree, lines) = parse("1 + 1");
        let semantics = Semantics::new(&tree, &lines);
        semantics.internal_compiler_error(tree.root(), "oh no");
    }
}