use crate::protocol::ast::*;
use crate::protocol::inputsource::*;
use crate::protocol::library;
use crate::protocol::parser::{symbol_table::*, type_table::*, LexedModule};

type Unit = ();
pub(crate) type VisitorResult = Result<Unit, ParseError2>;

pub(crate) struct Ctx<'p> {
    heap: &'p mut Heap,
    module: &'p LexedModule,
    symbols: &'p mut SymbolTable,
    types: &'p mut TypeTable,
}

/// Visitor is a generic trait that will fully walk the AST. The default
/// implementation of the visitors is to not recurse. The exception is the
/// top-level `visit_definition`, `visit_stmt` and `visit_expr` methods, which
/// call the appropriate visitor function.
pub(crate) trait Visitor2 {
    // Entry point
    fn visit_module(&mut self, ctx: &mut Ctx) -> VisitorResult {
        let mut def_index = 0;
        loop {
            let definition_id = {
                let root = &ctx.heap[ctx.module.root_id];
                if def_index >= root.definitions.len() {
                    return Ok(())
                }

                root.definitions[def_index]
            };

            self.visit_definition(ctx, definition_id)
        }
    }

    // Definitions
    // --- enum matching
    fn visit_definition(&mut self, ctx: &mut Ctx, id: DefinitionId) -> VisitorResult {
        match &ctx.heap[id] {
            Definition::Enum(def) => self.visit_enum_definition(ctx, def.this),
            Definition::Struct(def) => self.visit_struct_definition(ctx, def.this),
            Definition::Component(def) => self.visit_component_definition(ctx, def.this),
            Definition::Function(def) => self.visit_function_definition(ctx, def.this)
        }
    }

    // --- enum variant handling
    fn visit_enum_definition(&mut self, _ctx: &mut Ctx, id: EnumId) -> VisitorResult { Ok(()) }
    fn visit_struct_definition(&mut self, _ctx: &mut Ctx, id: StructId) -> VisitorResult { Ok(()) }
    fn visit_component_definition(&mut self, _ctx: &mut Ctx, id: ComponentId) -> VisitorResult { Ok(()) }
    fn visit_function_definition(&mut self, _ctx: &mut Ctx, id: FunctionId) -> VisitorResult { Ok(()) }

    // Statements
    // --- enum matching
    fn visit_stmt(&mut self, ctx: &mut Ctx, id: StatementId) -> VisitorResult {
        match &ctx.heap[id] {
            Statement::Block(stmt) => self.visit_block_stmt(ctx, stmt.this),
            Statement::Local(stmt) => self.visit_local_stmt(ctx, stmt.this),
            Statement::Skip(stmt) => self.visit_skip_stmt(ctx, stmt.this),
            Statement::Labeled(stmt) => self.visit_labeled_stmt(ctx, stmt.this),
            Statement::If(stmt) => self.visit_if_stmt(ctx, stmt.this),
            Statement::EndIf(_stmt) => Ok(()),
            Statement::While(stmt) => self.visit_while_stmt(ctx, stmt.this),
            Statement::EndWhile(_stmt) => Ok(()),
            Statement::Break(stmt) => self.visit_break_stmt(ctx, stmt.this),
            Statement::Continue(stmt) => self.visit_continue_stmt(ctx, stmt.this),
            Statement::Synchronous(stmt) => self.visit_synchronous_stmt(ctx, stmt.this),
            Statement::EndSynchronous(_stmt) => Ok(()),
            Statement::Return(stmt) => self.visit_return_stmt(ctx, stmt.this),
            Statement::Assert(stmt) => self.visit_assert_stmt(ctx, stmt.this),
            Statement::Goto(stmt) => self.visit_goto_stmt(ctx, stmt.this),
            Statement::New(stmt) => self.visit_new_stmt(ctx, stmt.this),
            Statement::Put(stmt) => self.visit_put_stmt(ctx, stmt.this),
            Statement::Expression(stmt) => self.visit_expr_stmt(ctx, stmt.this),
        }
    }

    fn visit_local_stmt(&mut self, ctx: &mut Ctx, id: LocalStatementId) -> VisitorResult {
        match &ctx.heap[id] {
            LocalStatement::Channel(stmt) => self.visit_local_channel_stmt(ctx, stmt.this),
            LocalStatement::Memory(stmt) => self.visit_local_memory_stmt(ctx, stmt.this),
        }
    }

    // --- enum variant handling
    fn visit_block_stmt(&mut self, _ctx: &mut Ctx, _id: BlockStatementId) -> VisitorResult { Ok(()) }
    fn visit_local_memory_stmt(&mut self, _ctx: &mut Ctx, _id: MemoryStatementId) -> VisitorResult { Ok(()) }
    fn visit_local_channel_stmt(&mut self, _ctx: &mut Ctx, _id: ChannelStatementId) -> VisitorResult { Ok(()) }
    fn visit_skip_stmt(&mut self, _ctx: &mut Ctx, _id: SkipStatementId) -> VisitorResult { Ok(()) }
    fn visit_labeled_stmt(&mut self, _ctx: &mut Ctx, _id: LabeledStatementId) -> VisitorResult { Ok(()) }
    fn visit_if_stmt(&mut self, _ctx: &mut Ctx, _id: IfStatementId) -> VisitorResult { Ok(()) }
    fn visit_while_stmt(&mut self, _ctx: &mut Ctx, _id: WhileStatementId) -> VisitorResult { Ok(()) }
    fn visit_break_stmt(&mut self, _ctx: &mut Ctx, _id: BreakStatementId) -> VisitorResult { Ok(()) }
    fn visit_continue_stmt(&mut self, _ctx: &mut Ctx, _id: ContinueStatementId) -> VisitorResult { Ok(()) }
    fn visit_synchronous_stmt(&mut self, _ctx: &mut Ctx, _id: SynchronousStatementId) -> VisitorResult { Ok(()) }
    fn visit_return_stmt(&mut self, _ctx: &mut Ctx, _id: ReturnStatementId) -> VisitorResult { Ok(()) }
    fn visit_assert_stmt(&mut self, _ctx: &mut Ctx, _id: AssertStatementId) -> VisitorResult { Ok(()) }
    fn visit_goto_stmt(&mut self, _ctx: &mut Ctx, _id: GotoStatementId) -> VisitorResult { Ok(()) }
    fn visit_new_stmt(&mut self, _ctx: &mut Ctx, _id: NewStatementId) -> VisitorResult { Ok(()) }
    fn visit_put_stmt(&mut self, _ctx: &mut Ctx, _id: PutStatementId) -> VisitorResult { Ok(()) }
    fn visit_expr_stmt(&mut self, _ctx: &mut Ctx, _id: ExpressionStatementId) -> VisitorResult { Ok(()) }

    // Expressions
    // --- enum matching
    fn visit_expr(&mut self, ctx: &mut Ctx, id: ExpressionId) -> VisitorResult {
        match &ctx.heap[id] {
            Expression::Assignment(expr) => self.visit_assignment_expr(ctx, expr.this),
            Expression::Conditional(expr) => self.visit_conditional_expr(ctx, expr.this),
            Expression::Binary(expr) => self.visit_binary_expr(ctx, expr.this),
            Expression::Unary(expr) => self.visit_unary_expr(ctx, expr.this),
            Expression::Indexing(expr) => self.visit_indexing_expr(ctx, expr.this),
            Expression::Slicing(expr) => self.visit_slicing_expr(ctx, expr.this),
            Expression::Select(expr) => self.visit_select_expr(ctx, expr.this),
            Expression::Array(expr) => self.visit_array_expr(ctx, expr.this),
            Expression::Constant(expr) => self.visit_constant_expr(ctx, expr.this),
            Expression::Call(expr) => self.visit_call_expr(ctx, expr.this),
            Expression::Variable(expr) => self.visit_variable_expr(ctx, expr.this),
        }
    }

    fn visit_assignment_expr(&mut self, _ctx: &mut Ctx, _id: AssignmentExpressionId) -> VisitorResult { Ok(()) }
    fn visit_conditional_expr(&mut self, _ctx: &mut Ctx, _id: ConditionalExpressionId) -> VisitorResult { Ok(()) }
    fn visit_binary_expr(&mut self, _ctx: &mut Ctx, _id: BinaryExpressionId) -> VisitorResult { Ok(()) }
    fn visit_unary_expr(&mut self, _ctx: &mut Ctx, _id: UnaryExpressionId) -> VisitorResult { Ok(()) }
    fn visit_indexing_expr(&mut self, _ctx: &mut Ctx, _id: IndexingExpressionId) -> VisitorResult { Ok(()) }
    fn visit_slicing_expr(&mut self, _ctx: &mut Ctx, _id: SlicingExpressionId) -> VisitorResult { Ok(()) }
    fn visit_select_expr(&mut self, _ctx: &mut Ctx, _id: SelectExpressionId) -> VisitorResult { Ok(()) }
    fn visit_array_expr(&mut self, _ctx: &mut Ctx, _id: ArrayExpressionId) -> VisitorResult { Ok(()) }
    fn visit_constant_expr(&mut self, _ctx: &mut Ctx, _id: ConstantExpressionId) -> VisitorResult { Ok(()) }
    fn visit_call_expr(&mut self, _ctx: &mut Ctx, _id: CallExpressionId) -> VisitorResult { Ok(()) }
    fn visit_variable_expr(&mut self, _ctx: &mut Ctx, _id: VariableExpressionId) -> VisitorResult { Ok(()) }
}

enum DefinitionType {
    Primitive,
    Composite,
    Function
}

/// This particular visitor will go through the entire AST in a recursive manner
/// and check if all statements and expressions are legal (e.g. no "return"
/// statements in component definitions), and will link certain AST nodes to
/// their appropriate targets (e.g. goto statements, or function calls).
///
/// This visitor will not perform control-flow analysis (e.g. making sure that
/// each function actually returns) and will also not perform type checking. So
/// the linking of function calls and component instantiations will be checked
/// and linked to the appropriate definitions, but the return types and/or
/// arguments will not be checked for validity.
///
/// The visitor visits each statement in a block in a breadth-first manner
/// first. We are thereby sure that we have found all variables/labels in a
/// particular block. In this phase nodes may queue statements for insertion
/// (e.g. the insertion of an `EndIf` statement for a particular `If`
/// statement). These will be inserted after visiting every node, after which
/// the visitor recurses into each statement in a block.
///
/// Because of this scheme expressions will not be visited in the breadth-first
/// pass.
struct ValidityAndLinkerVisitor {
    /// `in_sync` is `Some(id)` if the visitor is visiting the children of a
    /// synchronous statement. A single value is sufficient as nested
    /// synchronous statements are not allowed
    in_sync: Option<SynchronousStatementId>,
    /// `in_while` contains the last encountered `While` statement. This is used
    /// to resolve unlabeled `Continue`/`Break` statements.
    in_while: Option<WhileStatementId>,
    // Traversal state: current scope (which can be used to find the parent
    // scope), the definition variant we are considering, and whether the
    // visitor is performing breadthwise block statement traversal.
    cur_scope: Option<Scope>,
    def_type: DefinitionType,
    performing_breadth_pass: bool,
    // Keeping track of relative position in block in the breadth-first pass.
    // May not correspond to block.statement[index] if any statements are
    // inserted after the breadth-pass
    relative_pos_in_block: u32,
    // Single buffer of statement IDs that we want to traverse in a block.
    // Required to work around Rust borrowing rules
    // TODO: Maybe remove this in the future
    statement_buffer: Vec<StatementId>,
    // Statements to insert after the breadth pass in a single block
    insert_buffer: Vec<(u32, StatementId)>,
}

impl ValidityAndLinkerVisitor {
    fn new() -> Self {
        Self{
            in_sync: None,
            in_while: None,
            cur_scope: None,
            def_type: DefinitionType::Primitive,
            performing_breadth_pass: false,
            relative_pos_in_block: 0,
            statement_buffer: Vec::with_capacity(256),
            insert_buffer: Vec::with_capacity(32),
        }
    }

    fn reset_state(&mut self) {
        self.in_sync = None;
        self.in_while = None;
        self.cur_scope = None;
        self.def_type = DefinitionType::Primitive;
        self.relative_pos_in_block = 0;
        self.performing_breadth_pass = false;
        self.statement_buffer.clear();
        self.insert_buffer.clear();
    }
}

impl Visitor2 for ValidityAndLinkerVisitor {
    //--------------------------------------------------------------------------
    // Definition visitors
    //--------------------------------------------------------------------------

    fn visit_component_definition(&mut self, ctx: &mut Ctx, id: ComponentId) -> VisitorResult {
        self.reset_state();

        let block_id = {
            let def = &ctx.heap[id];
            match def.variant {
                ComponentVariant::Primitive => self.def_type = DefinitionType::Primitive,
                ComponentVariant::Composite => self.def_type = DefinitionType::Composite,
            }

            let body = ctx.heap[def.body].as_block_mut();

            self.statement_buffer.extend_from_slice(&body.statements);
            self.statement_stack_indices.push(0);
            body.this
        };

        self.cur_scope = Some(Scope {
            variant: ScopeVariant::Definition(id.upcast()),
        });

        self.performing_breadth_pass = true;
        self.visit_block_stmt(ctx, block_id)?;
        self.performing_breadth_pass = false;
        self.visit_block_stmt(ctx, block_id)
    }

    fn visit_function_definition(&mut self, ctx: &mut Ctx, id: FunctionId) -> VisitorResult {
        self.reset_state();

        // Set internal statement indices
        let block_id = {
            let def = &ctx.heap[id];
            self.def_type = DefinitionType::Function;
            let body = ctx.heap[def.body].as_block_mut();

            self.statement_buffer.extend_from_slice(&body.statements);
            self.statement_stack_indices.push(0);
            body.this
        };

        self.cur_scope = Some(Scope {
            variant: ScopeVariant::Definition(id.upcast()),
        });

        self.performing_breadth_pass = true;
        self.visit_block_stmt(ctx, block_id)?;
        self.performing_breadth_pass = false;
        self.visit_block_stmt(ctx, block_id)
    }

    //--------------------------------------------------------------------------
    // Statement visitors
    //--------------------------------------------------------------------------

    fn visit_block_stmt(&mut self, ctx: &mut Ctx, id: BlockStatementId) -> VisitorResult {
        self.visit_block_stmt_with_hint(ctx, id, None)
    }

    fn visit_local_memory_stmt(&mut self, ctx: &mut Ctx, id: MemoryStatementId) -> VisitorResult {
        if self.performing_breadth_pass {
            let stmt = &ctx.heap[id];
            stmt.relative_pos_in_block = self.relative_pos_in_block;
            self.checked_local_add(ctx, stmt.variable)?;
        } else {
            self.visit_expr(ctx, ctx.heap[id].initial)?;
        }

        Ok(())
    }

    fn visit_labeled_stmt(&mut self, ctx: &mut Ctx, id: LabeledStatementId) -> VisitorResult {
        if self.performing_breadth_pass {
            // Retrieve scope
            let scope = self.cur_scope.as_ref().unwrap();
            debug_assert!(scope.statement.is_some(), "expected scope statement at labeled stmt");
            debug_assert_eq!(
                scope.variant == ScopeVariant::Synchronous,
                self.in_sync.is_some(),
                "in synchronous scope variant, but 'in_sync' not set"
            );

            // Add label to block lookup
            self.checked_label_add(ctx, id)?;

            // Modify labeled statement itself
            let labeled = &mut ctx.heap[id];
            labeled.relative_pos_in_block = self.relative_pos_in_block;
            labeled.in_sync = if scope.variant == ScopeVariant::Synchronous {
                self.in_sync.clone()
            } else {
                None
            };
        }

        let body_id = ctx.heap[id].body;
        self.visit_stmt(ctx, body_id)?;

        Ok(())
    }

    fn visit_if_stmt(&mut self, ctx: &mut Ctx, id: IfStatementId) -> VisitorResult {
        if self.performing_breadth_pass {
            let position = ctx.heap[id].position;
            let end_if_id = ctx.heap.alloc_end_if_statement(|this| {
                EndIfStatement {
                    this,
                    start_if: id,
                    position,
                    next: None,
                }
            });
            let stmt = &mut ctx.heap[id];
            stmt.end_if = Some(end_if_id);
            self.insert_buffer.push((self.relative_pos_in_block + 1, end_if_id.upcast()));
        } else {
            // Traverse expression and bodies
            let (test_id, true_id, false_id) = {
                let stmt = &ctx.heap[id];
                (stmt.test, stmt.true_body, stmt.false_body)
            };
            self.visit_expr(ctx, test_id)?;
            self.visit_stmt(ctx, true_id)?;
            self.visit_stmt(ctx, false_id)?;
        }

        Ok(())
    }

    fn visit_while_stmt(&mut self, ctx: &mut Ctx, id: WhileStatementId) -> VisitorResult {
        if self.performing_breadth_pass {
            let scope = self.cur_scope.as_ref().unwrap();
            let position = ctx.heap[id].position;
            debug_assert_eq!(
                scope.variant == ScopeVariant::Synchronous,
                self.in_sync.is_some(),
                "in synchronous scope variant, but 'in_sync' not set"
            );
            let end_while_id = ctx.heap.alloc_end_while_statement(|this| {
                EndWhileStatement {
                    this,
                    start_while: Some(id),
                    position,
                    next: None,
                }
            });
            let stmt = &mut ctx.heap[id];
            stmt.end_while = Some(end_while_id);
            stmt.in_sync = self.in_sync.clone();

            self.insert_buffer.push((self.relative_pos_in_block + 1, end_while_id.upcast()));
        } else {
            let (test_id, body_id) = {
                let stmt = &ctx.heap[id];
                (stmt.test, stmt.body)
            };
            let old_while = self.in_while.replace(id);
            self.visit_expr(ctx, test_id)?;
            self.visit_stmt(ctx, body_id)?;
            self.in_while = old_while;
        }

        Ok(())
    }

    fn visit_break_stmt(&mut self, ctx: &mut Ctx, id: BreakStatementId) -> VisitorResult {
        if self.performing_breadth_pass {
            // Should be able to resolve break statements with a label in the
            // breadth pass, no need to do after resolving all labels
            let target_end_while = {
                let stmt = &ctx.heap[id];
                let target_while_id = self.resolve_break_or_continue_target(ctx, stmt.position, &stmt.label)?;
                let target_while = &ctx.heap[target_while_id];
                debug_assert!(target_while.end_while.is_some());
                target_while.end_while.unwrap()
            };

            let stmt = &mut ctx.heap[id];
            stmt.target = Some(target_end_while);
        }

        Ok(())
    }

    fn visit_continue_stmt(&mut self, ctx: &mut Ctx, id: ContinueStatementId) -> VisitorResult {
        if self.performing_breadth_pass {
            let target_while_id = {
                let stmt = &ctx.heap[id];
                self.resolve_break_or_continue_target(ctx, stmt.position, &stmt.label)?
            };

            let stmt = &mut ctx.heap[id];
            stmt.target = Some(target_while_id)
        }

        Ok(())
    }

    fn visit_synchronous_stmt(&mut self, ctx: &mut Ctx, id: SynchronousStatementId) -> VisitorResult {
        if self.performing_breadth_pass {
            // Check for validity of synchronous statement
            let cur_sync = &ctx.heap[id];
            if self.in_sync.is_some() {
                // Nested synchronous statement
                let old_sync = &ctx.heap[self.in_sync.unwrap()];
                return Err(
                    ParseError2::new_error(&ctx.module.source, cur_sync.position, "Illegal nested synchronous statement")
                        .with_postfixed_info(&ctx.module.source, old_sync.position, "It is nested in this synchronous statement")
                );
            }

            if self.def_type != DefinitionType::Primitive {
                return Err(ParseError2::new_error(
                    &ctx.module.source, cur_sync.position,
                    "Synchronous statements may only be used in primitive components"
                ));
            }

            // Append SynchronousEnd pseudo-statement
            let sync_end_id = ctx.heap.alloc_end_synchronous_statement(|this| EndSynchronousStatement{
                this,
                position: stmt.position,
                start_sync: id,
                next: None,
            });
            let sync_start = &mut ctx.heap[id];
            sync_start.end_sync = Some(sync_end_id);
            self.insert_buffer.push((self.relative_pos_in_block + 1, sync_end_id.upcast()));
        } else {
            let old = self.in_sync.replace(id);
            self.visit_stmt_with_hint(ctx, stmt.body, Some(id))?;
            self.in_sync = old;
        }

        Ok(())
    }

    fn visit_return_stmt(&mut self, ctx: &mut Ctx, id: ReturnStatementId) -> VisitorResult {
        if self.performing_breadth_pass {
            let stmt = &ctx.heap[id];
            if self.def_type != DefinitionType::Function {
                return Err(
                    ParseError2::new_error(&ctx.module.source, stmt.position, "Return statements may only appear in function bodies")
                );
            }
        } else {
            // If here then we are within a function
            self.visit_expr(ctx, ctx.heap[id].expression)?;
        }

        Ok(())
    }

    fn visit_assert_stmt(&mut self, ctx: &mut Ctx, id: AssertStatementId) -> VisitorResult {
        let stmt = &ctx.heap[id];
        if self.performing_breadth_pass {
            if self.def_type == DefinitionType::Function {
                // TODO: We probably want to allow this. Mark the function as
                //  using asserts, and then only allow calls to these functions
                //  within components. Such a marker will cascade through any
                //  functions that then call an asserting function
                return Err(
                    ParseError2::new_error(&ctx.module.source, stmt.position, "Illegal assert statement in a function")
                );
            }

            // We are in a component of some sort, but we also need to be within a
            // synchronous statement
            if self.in_sync.is_none() {
                return Err(
                    ParseError2::new_error(&ctx.module.source, stmt.position, "Illegal assert statement outside of a synchronous block")
                );
            }
        } else {
            self.visit_expr(ctx, stmt.expression)?;
        }

        Ok(())
    }

    fn visit_goto_stmt(&mut self, ctx: &mut Ctx, id: GotoStatementId) -> VisitorResult {
        if !self.performing_breadth_pass {
            // Must perform goto label resolving after the breadth pass, this
            // way we are able to find all the labels in current and outer
            // scopes.
            let goto_stmt = &mut ctx.heap[id];
            let target_id = self.find_label(ctx, &goto_stmt.label)?;
            goto_stmt.target = Some(target_id);

            let target = &ctx.heap[target_id];
            if self.in_sync != target.in_sync {
                // We can only goto the current scope or outer scopes. Because
                // nested sync statements are not allowed so if the value does
                // not match, then we must be inside a sync scope
                debug_assert!(self.in_sync.is_some());
                let sync_stmt = &ctx.heap[self.in_sync.unwrap()];
                return Err(
                    ParseError2::new_error(&ctx.module.source, goto_stmt.position, "Goto may not escape the surrounding synchronous block")
                        .with_postfixed_info(&ctx.module.source, target.position, "This is the target of the goto statement")
                        .with_postfixed_info(&ctx.module.source, sync_stmt.position, "Which will jump past this statement")
                );
            }
        }

        Ok(())
    }

    fn visit_new_stmt(&mut self, ctx: &mut Ctx, id: NewStatementId) -> VisitorResult {
        if self.performing_breadth_pass {
            // TODO: Cleanup error messages, can be done cleaner
            // Make sure new statement occurs within a composite component
            let new_stmt = &ctx.heap[id];
            if self.def_type != DefinitionType::Composite {
                return Err(
                    ParseError2::new_error(&ctx.module.source, new_stmt.position, "Instantiating components may only be done in composite components")
                );
            }

            // No fancy recursive parsing, must be followed by a call expression
            let definition_id = {
                let call_expr = &ctx.heap[new_stmt.expression];
                if let Expression::Call(call_expr) = call_expr {
                    if let Method::Symbolic(symbolic) = &call_expr.method {
                        // Resolve method
                        let (symbol, iter) = ctx.symbols.resolve_namespaced_symbol(ctx.module.root_id, &symbolic.identifier)?;
                        if iter.num_remaining() != 0 {
                            return Err(
                                ParseError2::new_error(&ctx.module.source, symbolic.identifier.position, "Unknown component")
                            )
                        }

                        match symbol.symbol {
                            Symbol::Namespace(_) => return Err(
                                ParseError2::new_error(&ctx.module.source, symbolic.identifier.position, "Unknown component")
                                    .with_postfixed_info(&ctx.module.source, symbol.position, "The identifier points to this import")
                            ),
                            Symbol::Definition((target_root_id, target_definition_id)) => {
                                match &ctx.heap[target_definition_id] {
                                    Definition::Component(_) => target_definition_id,
                                    _ => return Err(
                                        ParseError2::new_error(&ctx.module.source, symbolic.identifier.position, "Must instantiate a component")
                                    )
                                }
                            }
                        }
                    } else {
                        return Err(
                            ParseError2::new_error(&ctx.module.source, call_expr.position, "Must instantiate a component")
                        );
                    }
                } else {
                    return Err(
                        ParseError2::new_error(&ctx.module.source, call_expr.position, "Must instantiate a component")
                    );
                }
            };

            // Modify new statement's symbolic call to point to the appropriate
            // definition.
            let call_expr = &mut ctx.heap[new_stmt.expression];
            match &mut call_expr.method {
                Method::Symbolic(method) => method.definition = Some(definition_id),
                _ => unreachable!()
            }
        }

        Ok(())
    }

    fn visit_put_stmt(&mut self, ctx: &mut Ctx, id: PutStatementId) -> VisitorResult {
        // TODO: Make `put` an expression. Perhaps silly, but much easier to
        //  perform typechecking
        if self.performing_breadth_pass {
            let put_stmt = &ctx.heap[id];
            if self.in_sync.is_none() {
                return Err(ParseError2::new_error(
                    &ctx.module.source, put_stmt.position, "Put must be called in a synchronous block"
                ));
            }
        } else {
            let put_stmt = &ctx.heap[id];
            self.visit_expr(ctx, put_stmt.port)?;
            self.visit_expr(ctx, put_stmt.message)?;
        }

        Ok(())
    }

    fn visit_expr_stmt(&mut self, ctx: &mut Ctx, id: ExpressionStatementId) -> VisitorResult {
        if !self.performing_breadth_pass {
            let expr = &ctx.heap[id];
            self.visit_expr(ctx, expr.expression)?;
        }

        Ok(())
    }


    //--------------------------------------------------------------------------
    // Expression visitors
    //--------------------------------------------------------------------------

    fn visit_assignment_expr(&mut self, ctx: &mut Ctx, id: AssignmentExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);
        Ok(())
    }

    fn visit_conditional_expr(&mut self, ctx: &mut Ctx, id: ConditionalExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);
        let conditional_expr = &ctx.heap[id];
        self.visit_expr(ctx, conditional_expr.test)?;
        self.visit_expr(ctx, conditional_expr.true_expression)?;
        self.visit_expr(ctx, conditional_expr.false_expression)?;
        Ok(())
    }

    fn visit_binary_expr(&mut self, ctx: &mut Ctx, id: BinaryExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);
        let binary_expr = &ctx.heap[id];
        self.visit_expr(ctx, binary_expr.left)?;
        self.visit_expr(ctx, binary_expr.right)?;
        Ok(())
    }

    fn visit_unary_expr(&mut self, ctx: &mut Ctx, id: UnaryExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);
        let unary_expr = &ctx.heap[id];
        self.visit_expr(ctx, unary_expr.expression)?;
        Ok(())
    }

    fn visit_indexing_expr(&mut self, ctx: &mut Ctx, id: IndexingExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);
        let indexing_expr = &ctx.heap[id];
        self.visit_expr(ctx, indexing_expr.subject)?;
        self.visit_expr(ctx, indexing_expr.index)?;
        Ok(())
    }

    fn visit_slicing_expr(&mut self, ctx: &mut Ctx, id: SlicingExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);
        // TODO: Same as the select expression: slicing depends on the type of
        //  the thing that is being sliced.
        let slicing_expr = &ctx.heap[id];
        self.visit_expr(ctx, slicing_expr.subject)?;
        self.visit_expr(ctx, slicing_expr.from_index)?;
        self.visit_expr(ctx, slicing_expr.to_index)?;
        Ok(())
    }

    fn visit_select_expr(&mut self, ctx: &mut Ctx, id: SelectExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);
        // TODO: Is it true that this always depends on the return value? I
        //  mean: the following should be a valid expression:
        //  int i = some_call_that_returns_a_struct(5, 2).field_of_struct
        //  We could rule out some things (of which we're sure what the type is)
        //  but it seems better to do this later
        let select_expr = &ctx.heap[id];
        self.visit_expr(ctx, select_expr.subject)?;

        Ok(())
    }

    fn visit_array_expr(&mut self, ctx: &mut Ctx, id: ArrayExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);
        let array_expr = &ctx.heap[id];
        for field in &array_expr.elements {
            self.visit_expr(ctx, *field)?;
        }
        Ok(())
    }

    fn visit_constant_expr(&mut self, ctx: &mut Ctx, id: ConstantExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);
        Ok(())
    }

    fn visit_call_expr(&mut self, ctx: &mut Ctx, id: CallExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);

        let call_expr = &mut ctx.heap[id];
        match &mut call_expr.method {
            Method::Create => {},
            Method::Fires => {
                if self.def_type != DefinitionType::Primitive {
                    return Err(ParseError2::new_error(
                        &ctx.module.source, call_expr.position,
                        "A call to 'fires' may only occur in primitive component definitions"
                    ));
                }
            },
            Method::Get => {
                if self.def_type != DefinitionType::Primitive {
                    return Err(ParseError2::new_error(
                        &ctx.module.source, call_expr.position,
                        "A call to 'get' may only occur in primitive component definitions"
                    ));
                }
            },
            Method::Symbolic(symbolic) => {
                // Find symbolic method
                let symbol = ctx.symbols.resolve_namespaced_symbol(ctx.module.root_id, &symbolic.identifier);
                if symbol.is_none() {
                    return Err(ParseError2::new_error(
                        &ctx.module.source, symbolic.identifier.position,
                        "Could not find definition of function call"
                    ));
                }
                let (symbol, iter) = symbol.unwrap();
                if iter.num_remaining() != 0 {
                    return Err(
                        ParseError2::new_error(&ctx.module.source, symbolic.identifier.position,"Could not find definition of function call")
                            .with_postfixed_info(&ctx.module.source, symbol.position, "Could resolve part of the identifier to this symbol")
                    );
                }
                let definition_id = match &symbol.symbol {
                    Symbol::Definition((_, definition_id)) => {
                        let definition = ctx.types.get_definition(definition_id);
                        debug_assert!(definition.is_some(), "Symbol resolved to definition, but not present in type table");
                        let definition = definition.unwrap();
                        match definition {
                            DefinedType::Function(_) => Some(*definition_id),
                            _ => None,
                        }
                    },
                    Symbol::Namespace(_) => None,
                };
                if definition_id.is_none() {
                    return Err(
                        ParseError2::new_error(&ctx.module.source, symbolic.identifier.position, "Could not find definition of function call")
                    );
                }

                symbolic.definition = Some(definition_id.unwrap());
            }
        }

        Ok(())
    }

    fn visit_variable_expr(&mut self, ctx: &mut Ctx, id: VariableExpressionId) -> VisitorResult {
        debug_assert!(!self.performing_breadth_pass);

        let var_expr = &ctx.heap[id];
        let variable_id = self.find_variable(ctx, self.relative_pos_in_block, &var_expr.identifier)?;
        let var_expr = &mut ctx.heap[id];
        var_expr.declaration = Some(variable_id);

        Ok(())
    }
}

impl ValidityAndLinkerVisitor {
    //--------------------------------------------------------------------------
    // Special traversal
    //--------------------------------------------------------------------------

    fn visit_stmt_with_hint(&mut self, ctx: &mut Ctx, id: StatementId, hint: Option<SynchronousStatementId>) -> VisitorResult {
        if let Statement::Block(block) = &ctx.heap[id] {
            self.visit_block_stmt_with_hint(ctx, block.this, hint)
        } else {
            self.visit_stmt(ctx, id)
        }
    }

    fn visit_block_stmt_with_hint(&mut self, ctx: &mut Ctx, id: BlockStatementId, hint: Option<SynchronousStatementId>) -> VisitorResult {
        if self.performing_breadth_pass {
            // Our parent is performing a breadth-pass. We do this simple stuff
            // here
            let body = &mut ctx.heap[id];
            body.parent_scope = self.cur_scope.clone();
            body.relative_pos_in_parent = self.relative_pos_in_block;

            return Ok(())
        }

        // We may descend into children of this block. However, this is
        // where we first perform a breadth-first pass
        self.performing_breadth_pass = true;
        let old_scope = self.cur_scope.replace(match hint {
            Some(sync_id) => Scope{
                variant: ScopeVariant::Synchronous((sync_id, id)),
            },
            None => Scope{
                variant: ScopeVariant::Regular(id)
            }
        });
        let first_statement_index = self.statement_buffer.len();

        {
            let body = &ctx.heap[id];
            self.statement_buffer.extend_from_slice(&body.statements);
        }

        let mut stmt_index = first_statement_index;
        while stmt_index < self.statement_buffer.len() {
            self.relative_pos_in_block = (stmt_index - first_statement_index) as u32;
            self.visit_stmt(ctx, self.statement_buffer[stmt_index])?;
            stmt_index += 1;
        }

        if !self.insert_buffer.is_empty() {
            let body = &mut ctx.heap[id];
            for (pos, stmt) in self.insert_buffer.drain(..) {
                body.statements.insert(pos as usize, stmt);
            }
        }

        // And the depth pass. Because we're not actually visiting any inserted
        // nodes because we're using the statement buffer, we may safely use the
        // relative_pos_in_block counter.
        self.performing_breadth_pass = false;
        stmt_index = first_statement_index;
        while stmt_index < self.statement_buffer.len() {
            self.relative_pos_in_block = (stmt_index - first_statement_index) as u32;
            self.visit_stmt(ctx, self.statement_buffer[stmt_index])?;
            stmt_index += 1;
        }

        self.cur_scope = old_scope;

        // Pop statement buffer
        debug_assert!(self.insert_buffer.is_empty(), "insert buffer not empty after depth pass");
        self.statement_buffer.truncate(first_statement_index);

        Ok(())
    }

    //--------------------------------------------------------------------------
    // Utilities
    //--------------------------------------------------------------------------

    fn checked_local_add(&mut self, ctx: &mut Ctx, id: LocalId) -> Result<(), ParseError2> {
        debug_assert!(self.cur_scope.is_some());

        // Make sure we do not conflict with any global symbols
        {
            let ident = &ctx.heap[id].identifier;
            if let Some(symbol) = ctx.symbols.resolve_symbol(ctx.module.root_id, &ident.value) {
                return Err(
                    ParseError2::new_error(&ctx.module.source, ident.position, "Local variable declaration conflicts with symbol")
                        .with_postfixed_info(&ctx.module.source, symbol.position, "Conflicting symbol is found here")
                );
            }
        }

        // Make sure we do not shadow any variables in any of the scopes. Note
        // that variables in parent scopes may be declared later
        let local = &ctx.heap[id];
        let mut scope = self.cur_scope.as_ref().unwrap();
        let mut local_relative_pos = self.relative_pos_in_block;

        loop {
            debug_assert!(scope.is_block(), "scope is not a block");
            let block = &ctx.heap[scope.to_block()];
            for other_local_id in &block.locals {
                let other_local = &ctx.heap[*other_local_id];
                // Position check in case another variable with the same name
                // is defined in a higher-level scope, but later than the scope
                // in which the current variable resides.
                if local_relative_pos >= other_local.relative_pos_in_block && local.identifier.value == other_local.identifier.pos {
                    // Collision within this scope
                    return Err(
                        ParseError2::new_error(&ctx.module.source, local.position, "Local variable name conflicts with another variable")
                            .with_postfixed_info(&ctx.module.source, other_local.position, "Previous variable is found here")
                    );
                }
            }

            // Current scope is fine, move to parent scope if any
            debug_assert!(scope.parent.is_some(), "block scope does not have a parent");
            scope = scope.parent.as_ref().unwrap();
            if let ScopeVariant::Definition(definition_id) = scope.variant {
                // At outer scope, check parameters of function/component
                for parameter_id in ctx.heap[definition_id].parameters() {
                    let parameter = &ctx.heap[*parameter_id];
                    if local.identifier.value == parameter.identifier.value {
                        return Err(
                            ParseError2::new_error(&ctx.module.source, local.position, "Local variable name conflicts with parameter")
                                .with_postfixed_info(&ctx.module.source, parameter.position, "Parameter definition is found here")
                        );
                    }
                }

                break;
            }

            // If here, then we are dealing with a block-like parent block
            local_relative_pos = ctx.heap[scope.to_block()].relative_pos_in_parent;
        }

        // No collisions at all
        let block = &mut ctx.heap[self.cur_scope.as_ref().unwrap().to_block()];
        block.locals.push(id);

        Ok(())
    }

    /// Finds a variable in the visitor's scope that must appear before the
    /// specified relative position within that block.
    fn find_variable(&self, ctx: &Ctx, mut relative_pos: u32, identifier: &NamespacedIdentifier) -> Result<VariableId, ParseError2> {
        debug_assert!(self.cur_scope.is_some());
        debug_assert!(identifier.num_namespaces > 0);

        // TODO: Update once globals are possible as well
        if identifier.num_namespaces > 1 {
            todo!("Implement namespaced constant seeking")
        }

        // TODO: May still refer to an alias of a global symbol using a single
        //  identifier in the namespace.
        // No need to use iterator over namespaces if here
        let mut scope = self.cur_scope.as_ref().unwrap();
        loop {
            debug_assert!(scope.is_block());
            let block = &ctx.heap[scope.to_block()];
            for local_id in &block.locals {
                let local = &ctx.heap[*local_id];
                if local.relative_pos_in_block < relative_pos && local.identifier.value == identifier.value {
                    return Ok(local_id.upcast());
                }
            }

            debug_assert!(block.parent_scope.is_some());
            scope = block.parent_scope.as_ref().unwrap();
            if !scope.is_block() {
                // Definition scope, need to check arguments to definition
                match scope.variant {
                    ScopeVariant::Definition(definition_id) => {
                        let definition = &ctx.heap[definition_id];
                        for parameter_id in definition.parameters() {
                            let parameter = &ctx.heap[*parameter_id];
                            if parameter.identifier.value == identifier.value {
                                return Ok(parameter_id.upcast());
                            }
                        }
                    },
                    _ => unreachable!(),
                }

                // Variable could not be found
                return Err(ParseError2::new_error(
                    &ctx.module.source, identifier.position, "This variable is not declared"
                ));
            } else {
                relative_pos = block.relative_pos_in_parent;
            }
        }
    }

    fn checked_label_add(&mut self, ctx: &mut Ctx, id: LabeledStatementId) -> Result<(), ParseError2> {
        debug_assert!(self.cur_scope.is_some());

        // Make sure label is not defined within the current scope or any of the
        // parent scope.
        let label = &ctx.heap[id];
        let mut scope = self.cur_scope.as_ref().unwrap();

        loop {
            debug_assert!(scope.is_block(), "scope is not a block");
            let block = &ctx.heap[scope.to_block()];
            for other_label_id in &block.labels {
                let other_label = &ctx.heap[*other_label_id];
                if other_label.label.value == label.label.value {
                    // Collision
                    return Err(
                        ParseError2::new_error(&ctx.module.source, label.position, "Label name conflicts with another label")
                            .with_postfixed_info(&ctx.module.source, other_label.position, "Other label is found here")
                    );
                }
            }

            debug_assert!(block.parent_scope.is_some(), "block scope does not have a parent");
            scope = block.parent_scope.as_ref().unwrap();
            if !scope.is_block() {
                break;
            }
        }

        // No collisions
        let block = &mut ctx.heap[self.cur_scope.as_ref().unwrap().to_block()];
        block.labels.push(id);

        Ok(())
    }

    /// Finds a particular labeled statement by its identifier. Once found it
    /// will make sure that the target label does not skip over any variable
    /// declarations within the scope in which the label was found.
    fn find_label(&self, ctx: &Ctx, identifier: &Identifier) -> Result<LabeledStatementId, ParseError2> {
        debug_assert!(self.cur_scope.is_some());

        let mut scope = self.cur_scope.as_ref().unwrap();
        loop {
            debug_assert!(scope.is_block(), "scope is not a block");
            let relative_scope_pos = ctx.heap[scope.to_block()].relative_pos_in_parent;

            let block = &ctx.heap[scope.to_block()];
            for label_id in &block.labels {
                let label = &ctx.heap[*label_id];
                if label.label.value == identifier.value {
                    for local_id in &block.locals {
                        // TODO: Better to do this in control flow analysis, it
                        //  is legal to skip over a variable declaration if it
                        //  is not actually being used. I might be missing
                        //  something here when laying out the bytecode...
                        let local = &ctx.heap[*local_id];
                        if local.relative_pos_in_block > relative_scope_pos && local.relative_pos_in_block < label.relative_pos_in_block {
                            return Err(
                                ParseError2::new_error(&ctx.module.source, identifier.position, "This target label skips over a variable declaration")
                                    .with_postfixed_info(&ctx.module.source, label.position, "Because it jumps to this label")
                                    .with_postfixed_info(&ctx.module.source, local.position, "Which skips over this variable")
                            );
                        }
                    }
                    return Ok(*label_id);
                }
            }

            debug_assert!(scope.parent.is_some(), "block scope does not have a parent");
            scope = scope.parent.as_ref().unwrap();
            if !scope.is_block() {
                return Err(ParseError2::new_error(&ctx.module.source, identifier.position, "Could not find this label"));
            }

        }
    }

    /// This function will check if the provided while statement ID has a block
    /// statement that is one of our current parents.
    fn has_parent_while_scope(&self, ctx: &Ctx, id: WhileStatementId) -> bool {
        debug_assert!(self.cur_scope.is_some());
        let mut scope = self.cur_scope.as_ref().unwrap();
        let while_stmt = &ctx.heap[id];
        loop {
            debug_assert!(scope.is_block());
            let block = scope.to_block();
            if while_stmt.body == block.upcast() {
                return true;
            }

            debug_assert!(scope.parent.is_some(), "block scope does not have a parent");
            scope = scope.parent.as_ref().unwrap();
            if !scope.is_block() {
                return false;
            }
        }
    }

    /// This function should be called while dealing with break/continue
    /// statements. It will try to find the targeted while statement, using the
    /// target label if provided. If a valid target is found then the loop's
    /// ID will be returned, otherwise a parsing error is constructed.
    /// The provided input position should be the position of the break/continue
    /// statement.
    fn resolve_break_or_continue_target(&self, ctx: &Ctx, position: InputPosition, label: &Option<Identifier>) -> Result<WhileStatementId, ParseError2> {
        let target = match label {
            Some(label) => {
                let target_id = self.find_label(ctx, label)?;

                // Make sure break target is a while statement
                let target = &ctx.heap[target_id];
                if let Statement::While(target_stmt) = &target.body {
                    // Even though we have a target while statement, the break might not be
                    // present underneath this particular labeled while statement
                    if !self.has_parent_while_scope(ctx, target_stmt.this) {
                        ParseError2::new_error(&ctx.module.source, label.position, "Break statement is not nested under the target label's while statement")
                            .with_postfixed_info(&ctx.module.source, target.position, "The targeted label is found here");
                    }

                    target_stmt.this
                } else {
                    return Err(
                        ParseError2::new_error(&ctx.module.source, label.position, "Incorrect break target label, it must target a while loop")
                            .with_postfixed_info(&ctx.module.source, target.position, "The targeted label is found here")
                    );
                }
            },
            None => {
                // Use the enclosing while statement, the break must be
                // nested within that while statement
                if self.in_while.is_none() {
                    return Err(
                        ParseError2::new_error(&ctx.module.source, position, "Break statement is not nested under a while loop")
                    );
                }

                self.in_while.unwrap()
            }
        };

        // We have a valid target for the break statement. But we need to
        // make sure we will not break out of a synchronous block
        {
            let target_while = &ctx.heap[target];
            if target_while.in_sync != self.in_sync {
                // Break is nested under while statement, so can only escape a
                // sync block if the sync is nested inside the while statement.
                debug_assert!(self.in_sync.is_some());
                let sync_stmt = &ctx.heap[self.in_sync.unwrap()];
                return Err(
                    ParseError2::new_error(&ctx.module.source, position, "Break may not escape the surrounding synchronous block")
                        .with_postfixed_info(&ctx.module.source, target_while.position, "The break escapes out of this loop")
                        .with_postfixed_info(&ctx.module.source, sync_stmt.position, "And would therefore escape this synchronous block")
                );
            }
        }

        Ok(target)
    }
}