handle variable scopes correctly, add option to write name section

src/ast.rs

@@ -58,6 +58,69 @@ pub struct Function {
     pub params: Vec<(String, Type)>,
     pub type_: Option<Type>,
     pub body: Expression,
+    pub locals: Locals,
+}
+
+#[derive(Debug, Default)]
+pub struct Locals {
+    pub params: Vec<Local>,
+    pub locals: Vec<Local>,
+}
+
+impl Locals {
+    pub fn add_param(&mut self, span: Span, name: String, type_: Type) -> u32 {
+        assert!(self.locals.is_empty());
+        let id = self.params.len() as u32;
+        self.params.push(Local {
+            span,
+            name,
+            type_,
+            index: Some(id),
+        });
+        id
+    }
+
+    pub fn add_local(&mut self, span: Span, name: String, type_: Type, store: bool) -> u32 {
+        let id = (self.params.len() + self.locals.len()) as u32;
+        self.locals.push(Local {
+            span,
+            name,
+            type_,
+            index: store.then(|| id),
+        });
+        id
+    }
+}
+
+impl std::ops::Index<u32> for Locals {
+    type Output = Local;
+    fn index(&self, id: u32) -> &Local {
+        let id = id as usize;
+        if id < self.params.len() {
+            &self.params[id]
+        } else {
+            &self.locals[id - self.params.len()]
+        }
+    }
+}
+
+impl std::ops::IndexMut<u32> for Locals {
+    fn index_mut(&mut self, id: u32) -> &mut Local {
+        let id = id as usize;
+        if id < self.params.len() {
+            &mut self.params[id]
+        } else {
+            &mut self.locals[id - self.params.len()]
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct Local {
+    pub span: Span,
+    pub name: String,
+    pub type_: Type,
+    pub index: Option<u32>,
 }
 
 #[derive(Debug)]
@@ -75,8 +138,8 @@ pub enum DataValues {
     String(String),
     File {
         path: PathBuf,
-        data: Vec<u8>
-    }
+        data: Vec<u8>,
+    },
 }
 
 #[derive(Debug, Clone)]
@@ -108,28 +171,28 @@ impl Expression {
     pub fn const_i32(&self) -> i32 {
         match self.expr {
             Expr::I32Const(v) => v,
-            _ => panic!("Expected I32Const")
+            _ => panic!("Expected I32Const"),
         }
     }
 
     pub fn const_i64(&self) -> i64 {
         match self.expr {
             Expr::I64Const(v) => v,
-            _ => panic!("Expected I64Const")
+            _ => panic!("Expected I64Const"),
         }
     }
 
     pub fn const_f32(&self) -> f32 {
         match self.expr {
             Expr::F32Const(v) => v,
-            _ => panic!("Expected F32Const")
+            _ => panic!("Expected F32Const"),
         }
     }
 
     pub fn const_f64(&self) -> f64 {
         match self.expr {
             Expr::F64Const(v) => v,
-            _ => panic!("Expected F64Const")
+            _ => panic!("Expected F64Const"),
         }
     }
 }
@@ -144,12 +207,16 @@ pub enum Expr {
     I64Const(i64),
     F32Const(f32),
     F64Const(f64),
-    Variable(String),
+    Variable {
+        name: String,
+        local_id: Option<u32>,
+    },
     Let {
         name: String,
         type_: Option<Type>,
         value: Option<Box<Expression>>,
         let_type: LetType,
+        local_id: Option<u32>,
     },
     Poke {
         mem_location: MemoryLocation,
@@ -181,10 +248,12 @@ pub enum Expr {
     Assign {
         name: String,
         value: Box<Expression>,
+        local_id: Option<u32>,
     },
     LocalTee {
         name: String,
         value: Box<Expression>,
+        local_id: Option<u32>,
     },
     Cast {
         value: Box<Expression>,
@@ -224,7 +293,7 @@ impl Expr {
     }
 }
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum LetType {
     Normal,
     Lazy,

src/emit.rs

@@ -2,13 +2,13 @@ use std::collections::HashMap;
 
 use wasm_encoder::{
     BlockType, CodeSection, DataSection, EntityType, Export, ExportSection, Function,
-    FunctionSection, GlobalSection, GlobalType, ImportSection, Instruction, MemArg, MemoryType,
-    Module, StartSection, TypeSection, ValType,
+    FunctionSection, GlobalSection, GlobalType, ImportSection, IndirectNameMap, Instruction,
+    MemArg, MemoryType, Module, NameMap, NameSection, StartSection, TypeSection, ValType,
 };
 
-use crate::{ast, intrinsics::Intrinsics};
+use crate::{ast, intrinsics::Intrinsics, Options};
 
-pub fn emit(script: &ast::Script) -> Vec<u8> {
+pub fn emit(script: &ast::Script, module_name: &str, options: &Options) -> Vec<u8> {
     let mut module = Module::new();
 
     let function_types = collect_function_types(script);
@@ -188,6 +188,56 @@ pub fn emit(script: &ast::Script) -> Vec<u8> {
         module.section(&data_section);
     }
 
+    if options.debug {
+        let mut names = NameSection::new();
+
+        names.module(module_name);
+
+        let mut functions = HashMap::new();
+        for (name, index) in &function_map {
+            functions.insert(*index, name);
+        }
+        let mut keys: Vec<_> = functions.keys().collect();
+        keys.sort();
+
+        let mut function_names = NameMap::new();
+        for i in keys {
+            function_names.append(*i, functions[i]);
+        }
+        names.functions(&function_names);
+
+        let mut functions = HashMap::new();
+        for function in &script.functions {
+            let mut local_names = NameMap::new();
+            for param in &function.locals.params {
+                local_names.append(param.index.unwrap(), &param.name);
+            }
+
+            let mut locals = HashMap::new();
+            for local in &function.locals.locals {
+                if let Some(index) = local.index {
+                    locals.insert(index, &local.name);
+                }
+            }
+            let mut keys: Vec<_> = locals.keys().collect();
+            keys.sort();
+            for i in keys {
+                local_names.append(*i, locals[i]);
+            }
+            functions.insert(*function_map.get(&function.name).unwrap(), local_names);
+        }
+        let mut keys: Vec<_> = functions.keys().collect();
+        keys.sort();
+
+        let mut locals = IndirectNameMap::new();
+        for i in keys {
+            locals.append(*i, &functions[i]);
+        }
+        names.locals(&locals);
+
+        module.section(&names);
+    }
+
     module.finish()
 }
 
@@ -237,9 +287,9 @@ struct FunctionContext<'a> {
     function: &'a mut Function,
     globals: &'a HashMap<&'a str, u32>,
     functions: &'a HashMap<String, u32>,
-    locals: &'a HashMap<String, u32>,
+    locals: &'a ast::Locals,
     labels: Vec<String>,
-    let_values: HashMap<&'a str, (&'a ast::Expression, ast::LetType)>,
+    let_values: HashMap<u32, (&'a ast::Expression, ast::LetType)>,
     intrinsics: &'a Intrinsics,
 }
 
@@ -249,27 +299,22 @@ fn emit_function(
     functions: &HashMap<String, u32>,
     intrinsics: &Intrinsics,
 ) -> Function {
-    let mut locals = Vec::new();
-    collect_locals_expr(&func.body, &mut locals);
-    locals.sort_by_key(|(_, t)| *t);
-
-    let mut function = Function::new_with_locals_types(locals.iter().map(|(_, t)| map_type(*t)));
-
-    let mut local_map: HashMap<String, u32> = HashMap::new();
-
-    for (ref name, _) in func.params.iter() {
-        local_map.insert(name.clone(), local_map.len() as u32);
-    }
-
-    for (name, _) in locals {
-        local_map.insert(name, local_map.len() as u32);
-    }
+    let mut function = Function::new_with_locals_types({
+        let mut locals: Vec<(u32, ast::Type)> = func
+            .locals
+            .locals
+            .iter()
+            .filter_map(|local| local.index.map(|i| (i, local.type_)))
+            .collect();
+        locals.sort();
+        locals.into_iter().map(|(_, t)| map_type(t))
+    });
 
     let mut context = FunctionContext {
         function: &mut function,
         globals,
         functions,
-        locals: &local_map,
+        locals: &func.locals,
         labels: vec![],
         let_values: HashMap::new(),
         intrinsics,
@@ -284,89 +329,6 @@ fn emit_function(
     function
 }
 
-fn collect_locals_expr<'a>(expr: &ast::Expression, locals: &mut Vec<(String, ast::Type)>) {
-    match &expr.expr {
-        ast::Expr::Block {
-            statements,
-            final_expression,
-        } => {
-            for stmt in statements {
-                collect_locals_expr(stmt, locals);
-            }
-            if let Some(ref expr) = final_expression {
-                collect_locals_expr(expr, locals);
-            }
-        }
-        ast::Expr::Let {
-            name, type_, value, ..
-        } => {
-            locals.push((name.clone(), type_.unwrap()));
-            if let Some(ref value) = value {
-                collect_locals_expr(value, locals);
-            }
-        }
-        ast::Expr::Peek(mem_location) => collect_locals_expr(&mem_location.left, locals),
-        ast::Expr::Poke {
-            mem_location,
-            value,
-            ..
-        } => {
-            collect_locals_expr(&mem_location.left, locals);
-            collect_locals_expr(value, locals);
-        }
-        ast::Expr::Variable { .. }
-        | ast::Expr::I32Const(_)
-        | ast::Expr::I64Const(_)
-        | ast::Expr::F32Const(_)
-        | ast::Expr::F64Const(_) => (),
-        ast::Expr::UnaryOp { value, .. } => collect_locals_expr(value, locals),
-        ast::Expr::BinOp { left, right, .. } => {
-            collect_locals_expr(left, locals);
-            collect_locals_expr(right, locals);
-        }
-        ast::Expr::Branch(_) => (),
-        ast::Expr::BranchIf { condition, .. } => collect_locals_expr(condition, locals),
-        ast::Expr::Assign { value, .. } => collect_locals_expr(value, locals),
-        ast::Expr::LocalTee { value, .. } => collect_locals_expr(value, locals),
-        ast::Expr::Loop { block, .. } => collect_locals_expr(block, locals),
-        ast::Expr::LabelBlock { block, .. } => collect_locals_expr(block, locals),
-        ast::Expr::Cast { value, .. } => collect_locals_expr(value, locals),
-        ast::Expr::FuncCall { params, .. } => {
-            for param in params {
-                collect_locals_expr(param, locals);
-            }
-        }
-        ast::Expr::Select {
-            condition,
-            if_true,
-            if_false,
-            ..
-        } => {
-            collect_locals_expr(condition, locals);
-            collect_locals_expr(if_true, locals);
-            collect_locals_expr(if_false, locals);
-        }
-        ast::Expr::If {
-            condition,
-            if_true,
-            if_false,
-        } => {
-            collect_locals_expr(condition, locals);
-            collect_locals_expr(if_true, locals);
-            if let Some(if_false) = if_false {
-                collect_locals_expr(if_false, locals);
-            }
-        }
-        ast::Expr::Return { value: Some(value) } => collect_locals_expr(value, locals),
-        ast::Expr::Return { value: None } => (),
-        ast::Expr::First { value, drop } => {
-            collect_locals_expr(value, locals);
-            collect_locals_expr(drop, locals);
-        }
-        ast::Expr::Error => unreachable!(),
-    }
-}
-
 fn mem_arg_for_location(mem_location: &ast::MemoryLocation) -> MemArg {
     let offset = if let ast::Expr::I32Const(v) = mem_location.right.expr {
         v as u32 as u64
@@ -405,19 +367,20 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
         }
         ast::Expr::Let {
             value,
-            name,
             let_type,
+            local_id,
             ..
         } => {
+            let local = &ctx.locals[local_id.unwrap()];
             if let Some(ref value) = value {
                 match let_type {
                     ast::LetType::Normal => {
                         emit_expression(ctx, value);
                         ctx.function
-                            .instruction(&Instruction::LocalSet(*ctx.locals.get(name).unwrap()));
+                            .instruction(&Instruction::LocalSet(local.index.unwrap()));
                     }
                     ast::LetType::Lazy | ast::LetType::Inline => {
-                        ctx.let_values.insert(name, (value, *let_type));
+                        ctx.let_values.insert(local_id.unwrap(), (value, *let_type));
                     }
                 }
             }
@@ -599,11 +562,16 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
         ast::Expr::F64Const(v) => {
             ctx.function.instruction(&Instruction::F64Const(*v));
         }
-        ast::Expr::Assign { name, value, .. } => {
+        ast::Expr::Assign {
+            name,
+            value,
+            local_id,
+            ..
+        } => {
             emit_expression(ctx, value);
-            if let Some(local_index) = ctx.locals.get(name) {
+            if let &Some(id) = local_id {
                 ctx.function
-                    .instruction(&Instruction::LocalSet(*local_index));
+                    .instruction(&Instruction::LocalSet(ctx.locals[id].index.unwrap()));
             } else if let Some(global_index) = ctx.globals.get(name.as_str()) {
                 ctx.function
                     .instruction(&Instruction::GlobalSet(*global_index));
@@ -611,10 +579,13 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
                 unreachable!();
             }
         }
-        ast::Expr::LocalTee { name, value, .. } => {
+        ast::Expr::LocalTee {
+            value, local_id, ..
+        } => {
             emit_expression(ctx, value);
-            let index = ctx.locals.get(name).unwrap();
-            ctx.function.instruction(&Instruction::LocalTee(*index));
+            ctx.function.instruction(&Instruction::LocalTee(
+                ctx.locals[local_id.unwrap()].index.unwrap(),
+            ));
         }
         ast::Expr::Loop { label, block, .. } => {
             ctx.labels.push(label.to_string());
@@ -632,14 +603,15 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
             ctx.labels.pop();
             ctx.function.instruction(&Instruction::End);
         }
-        ast::Expr::Variable(name) => {
-            if let Some(index) = ctx.locals.get(name) {
-                if let Some((expr, let_type)) = ctx.let_values.get(name.as_str()) {
+        ast::Expr::Variable { name, local_id } => {
+            if let &Some(id) = local_id {
+                if let Some((expr, let_type)) = ctx.let_values.get(&id) {
                     match let_type {
                         ast::LetType::Lazy => {
-                            let expr = ctx.let_values.remove(name.as_str()).unwrap().0;
+                            let expr = ctx.let_values.remove(&id).unwrap().0;
                             emit_expression(ctx, expr);
-                            ctx.function.instruction(&Instruction::LocalTee(*index));
+                            ctx.function
+                                .instruction(&Instruction::LocalTee(ctx.locals[id].index.unwrap()));
                         }
                         ast::LetType::Inline => {
                             let expr = *expr;
@@ -648,7 +620,8 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
                         _ => unreachable!(),
                     }
                 } else {
-                    ctx.function.instruction(&Instruction::LocalGet(*index));
+                    ctx.function
+                        .instruction(&Instruction::LocalGet(ctx.locals[id].index.unwrap()));
                 }
             } else if let Some(index) = ctx.globals.get(name.as_str()) {
                 ctx.function.instruction(&Instruction::GlobalGet(*index));

src/lib.rs

@@ -1,26 +1,41 @@
 use anyhow::{bail, Result};
+use std::ffi::OsStr;
 use std::io::prelude::*;
 use std::{fs::File, path::Path};
 
 mod ast;
 mod constfold;
 mod emit;
+mod includes;
 mod intrinsics;
 mod parser;
 mod typecheck;
-mod includes;
 
 type Span = std::ops::Range<usize>;
 
-pub fn compile_file<P: AsRef<Path>>(path: P) -> Result<Vec<u8>> {
+#[derive(Default)]
+pub struct Options {
+    pub(crate) debug: bool,
+}
+
+impl Options {
+    pub fn with_debug(self) -> Self {
+        Options {
+            debug: true,
+            ..self
+        }
+    }
+}
+
+pub fn compile_file<P: AsRef<Path>>(path: P, options: Options) -> Result<Vec<u8>> {
     let path = path.as_ref();
     let mut input = String::new();
     File::open(path)?.read_to_string(&mut input)?;
 
-    compile_str(&input, path)
+    compile_str(&input, path, options)
 }
 
-pub fn compile_str(input: &str, path: &Path) -> Result<Vec<u8>> {
+pub fn compile_str(input: &str, path: &Path, options: Options) -> Result<Vec<u8>> {
     let mut script = match parser::parse(&input) {
         Ok(script) => script,
         Err(_) => bail!("Parse failed"),
@@ -32,6 +47,13 @@ pub fn compile_str(input: &str, path: &Path) -> Result<Vec<u8>> {
     if let Err(_) = typecheck::tc_script(&mut script, &input) {
         bail!("Type check failed");
     }
-    let wasm = emit::emit(&script);
+    let wasm = emit::emit(
+        &script,
+        &path
+            .file_stem()
+            .unwrap_or_else(|| OsStr::new("unknown"))
+            .to_string_lossy(),
+        &options,
+    );
     Ok(wasm)
 }

src/main.rs

@@ -1,17 +1,21 @@
-use anyhow::{anyhow, Result};
+use anyhow::Result;
 use std::io::prelude::*;
 use std::{fs::File, path::PathBuf};
 
-use curlywas::compile_file;
+use curlywas::{compile_file, Options};
 
 fn main() -> Result<()> {
-    let mut filename = PathBuf::from(
-        std::env::args()
-            .nth(1)
-            .ok_or_else(|| anyhow!("Path to .hw file missing"))?,
-    );
+    let mut args = pico_args::Arguments::from_env();
 
-    let wasm = compile_file(&filename)?;
+    let mut options = Options::default();
+
+    if args.contains(["-d", "--debug"]) {
+        options = options.with_debug();
+    }
+
+    let mut filename = args.free_from_os_str::<PathBuf, bool>(|s| Ok(s.into()))?;
+
+    let wasm = compile_file(&filename, options)?;
 
     wasmparser::validate(&wasm)?;
 

src/parser.rs

@@ -309,7 +309,10 @@ fn script_parser() -> impl Parser<Token, ast::Script, Error = Simple<Token>> + C
             .labelled("value");
 
             let variable = filter_map(|span, tok| match tok {
-                Token::Ident(id) => Ok(ast::Expr::Variable(id)),
+                Token::Ident(id) => Ok(ast::Expr::Variable {
+                    name: id,
+                    local_id: None,
+                }),
                 _ => Err(Simple::expected_input_found(span, Vec::new(), Some(tok))),
             })
             .labelled("variable");
@@ -319,6 +322,7 @@ fn script_parser() -> impl Parser<Token, ast::Script, Error = Simple<Token>> + C
                 .map(|(name, expr)| ast::Expr::LocalTee {
                     name,
                     value: Box::new(expr),
+                    local_id: None,
                 })
                 .boxed();
 
@@ -385,6 +389,7 @@ fn script_parser() -> impl Parser<Token, ast::Script, Error = Simple<Token>> + C
                     type_,
                     value: value.map(Box::new),
                     let_type: let_type.unwrap_or(ast::LetType::Normal),
+                    local_id: None,
                 })
                 .boxed();
 
@@ -395,6 +400,7 @@ fn script_parser() -> impl Parser<Token, ast::Script, Error = Simple<Token>> + C
                 .map(|(name, value)| ast::Expr::Assign {
                     name,
                     value: Box::new(value),
+                    local_id: None,
                 })
                 .boxed();
 
@@ -807,6 +813,7 @@ fn script_parser() -> impl Parser<Token, ast::Script, Error = Simple<Token>> + C
                     name,
                     type_,
                     body,
+                    locals: ast::Locals::default(),
                 })
             })
             .boxed();

src/typecheck.rs

@@ -20,7 +20,8 @@ pub fn tc_script(script: &mut ast::Script, source: &str) -> Result<()> {
         source,
         global_vars: HashMap::new(),
         functions: HashMap::new(),
-        local_vars: HashMap::new(),
+        locals: ast::Locals::default(),
+        local_vars: LocalVars::new(),
         block_stack: Vec::new(),
         return_type: None,
         intrinsics: Intrinsics::new(),
@@ -122,22 +123,22 @@ pub fn tc_script(script: &mut ast::Script, source: &str) -> Result<()> {
 
     for f in &mut script.functions {
         context.local_vars.clear();
+        context.local_vars.push_scope();
         for (name, type_) in &f.params {
-            if let Some(Var { span, .. }) = context
+            if let Some(span) = context
                 .local_vars
                 .get(name)
-                .or_else(|| context.global_vars.get(name))
+                .map(|id| &context.locals[id].span)
+                .or_else(|| context.global_vars.get(name).map(|v| &v.span))
             {
                 result =
                     report_duplicate_definition("Variable already defined", &f.span, span, source);
             } else {
                 context.local_vars.insert(
                     name.clone(),
-                    Var {
-                        type_: *type_,
-                        span: f.span.clone(),
-                        mutable: true,
-                    },
+                    context
+                        .locals
+                        .add_param(f.span.clone(), name.clone(), *type_),
                 );
             }
         }
@@ -145,6 +146,22 @@ pub fn tc_script(script: &mut ast::Script, source: &str) -> Result<()> {
 
         tc_expression(&mut context, &mut f.body)?;
 
+        let mut local_mapping: Vec<(ast::Type, usize)> = context
+            .locals
+            .locals
+            .iter()
+            .enumerate()
+            .filter(|(_, local)| local.index.is_some())
+            .map(|(index, local)| (local.type_, index))
+            .collect();
+        local_mapping.sort_by_key(|&(t, _)| t);
+        let locals_start = context.locals.params.len();
+        for (id, (_, index)) in local_mapping.into_iter().enumerate() {
+            context.locals.locals[index].index = Some((locals_start + id) as u32);
+        }
+
+        f.locals = std::mem::replace(&mut context.locals, ast::Locals::default());
+
         if f.body.type_ != f.type_ {
             result = type_mismatch(f.type_, &f.span, f.body.type_, &f.body.span, source);
         }
@@ -233,12 +250,50 @@ struct Context<'a> {
     source: &'a str,
     global_vars: Vars,
     functions: HashMap<String, FunctionType>,
-    local_vars: Vars,
+    locals: ast::Locals,
+    local_vars: LocalVars,
     block_stack: Vec<String>,
     return_type: Option<ast::Type>,
     intrinsics: Intrinsics,
 }
 
+struct LocalVars(Vec<HashMap<String, u32>>);
+
+impl LocalVars {
+    fn new() -> LocalVars {
+        LocalVars(Vec::new())
+    }
+
+    fn get(&self, name: &str) -> Option<u32> {
+        self.0
+            .iter()
+            .rev()
+            .filter_map(|scope| scope.get(name))
+            .next()
+            .copied()
+    }
+
+    fn get_in_current(&self, name: &str) -> Option<u32> {
+        self.0.last().unwrap().get(name).copied()
+    }
+
+    fn clear(&mut self) {
+        self.0.clear();
+    }
+
+    fn push_scope(&mut self) {
+        self.0.push(HashMap::new());
+    }
+
+    fn pop_scope(&mut self) {
+        self.0.pop();
+    }
+
+    fn insert(&mut self, name: String, id: u32) {
+        self.0.last_mut().unwrap().insert(name, id);
+    }
+}
+
 fn report_duplicate_definition(
     msg: &str,
     span: &Span,
@@ -360,20 +415,25 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
             ref mut statements,
             ref mut final_expression,
         } => {
+            context.local_vars.push_scope();
             for stmt in statements {
                 tc_expression(context, stmt)?;
             }
-            if let Some(final_expression) = final_expression {
+            let type_ = if let Some(final_expression) = final_expression {
                 tc_expression(context, final_expression)?;
                 final_expression.type_
             } else {
                 None
-            }
+            };
+            context.local_vars.pop_scope();
+            type_
         }
         ast::Expr::Let {
             ref mut value,
             ref mut type_,
             ref name,
+            let_type,
+            ref mut local_id,
             ..
         } => {
             if let Some(ref mut value) = value {
@@ -395,26 +455,21 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
                 }
             }
             if let Some(type_) = type_ {
-                if let Some(Var { span, .. }) = context
+                let store = let_type != ast::LetType::Inline;
+                let id = context
                     .local_vars
-                    .get(name)
-                    .or_else(|| context.global_vars.get(name))
-                {
-                    return report_duplicate_definition(
-                        "Variable already defined",
-                        &expr.span,
-                        span,
-                        context.source,
-                    );
-                }
-                context.local_vars.insert(
-                    name.clone(),
-                    Var {
-                        type_: *type_,
-                        span: expr.span.clone(),
-                        mutable: true,
-                    },
-                );
+                    .get_in_current(name)
+                    .filter(|id| {
+                        let local = &context.locals[*id];
+                        local.type_ == *type_ && store == local.index.is_some()
+                    })
+                    .unwrap_or_else(|| {
+                        context
+                            .locals
+                            .add_local(expr.span.clone(), name.clone(), *type_, store)
+                    });
+                *local_id = Some(id);
+                context.local_vars.insert(name.clone(), id);
             } else {
                 Report::build(ReportKind::Error, (), expr.span.start)
                     .with_message("Type missing")
@@ -528,12 +583,14 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
                 }
             }
         }
-        ast::Expr::Variable(ref name) => {
-            if let Some(&Var { type_, .. }) = context
-                .global_vars
-                .get(name)
-                .or_else(|| context.local_vars.get(name))
-            {
+        ast::Expr::Variable {
+            ref name,
+            ref mut local_id,
+        } => {
+            if let Some(id) = context.local_vars.get(name) {
+                *local_id = Some(id);
+                Some(context.locals[id].type_)
+            } else if let Some(&Var { type_, .. }) = context.global_vars.get(name) {
                 Some(type_)
             } else {
                 return unknown_variable(&expr.span, context.source);
@@ -542,58 +599,61 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
         ast::Expr::Assign {
             ref name,
             ref mut value,
+            ref mut local_id,
         } => {
             tc_expression(context, value)?;
-            if let Some(&Var {
+
+            let (type_, span) = if let Some(id) = context.local_vars.get(name) {
+                *local_id = Some(id);
+                let local = &context.locals[id];
+                if local.index.is_none() {
+                    return immutable_assign(&expr.span, context.source);
+                }
+                (local.type_, &local.span)
+            } else if let Some(&Var {
                 type_,
                 ref span,
                 mutable,
-            }) = context
-                .local_vars
-                .get(name)
-                .or_else(|| context.global_vars.get(name))
+            }) = context.global_vars.get(name)
             {
-                if value.type_ != Some(type_) {
-                    return type_mismatch(
-                        Some(type_),
-                        span,
-                        value.type_,
-                        &value.span,
-                        context.source,
-                    );
-                }
                 if !mutable {
                     return immutable_assign(&expr.span, context.source);
                 }
+                (type_, span)
             } else {
                 return unknown_variable(&expr.span, context.source);
+            };
+
+            if value.type_ != Some(type_) {
+                return type_mismatch(Some(type_), span, value.type_, &value.span, context.source);
             }
             None
         }
         ast::Expr::LocalTee {
             ref name,
             ref mut value,
+            ref mut local_id,
         } => {
             tc_expression(context, value)?;
-            if let Some(&Var {
-                type_,
-                ref span,
-                mutable,
-            }) = context.local_vars.get(name)
-            {
-                if value.type_ != Some(type_) {
+            if let Some(id) = context.local_vars.get(name) {
+                *local_id = Some(id);
+                let local = &context.locals[id];
+
+                if local.index.is_none() {
+                    return immutable_assign(&expr.span, context.source);
+                }
+
+                if value.type_ != Some(local.type_) {
                     return type_mismatch(
-                        Some(type_),
-                        span,
+                        Some(local.type_),
+                        &local.span,
                         value.type_,
                         &value.span,
                         context.source,
                     );
                 }
-                if !mutable {
-                    return immutable_assign(&expr.span, context.source);
-                }
-                Some(type_)
+
+                Some(local.type_)
             } else {
                 return unknown_variable(&expr.span, context.source);
             }