can compile uw8loader

.gitignore

@@ -1,2 +1,4 @@
 /target
-*.uw8
+.cargo/
+*.wasm
+*.lua*

src/ast.rs

@@ -88,6 +88,7 @@ pub enum Expr {
         label: String,
         block: Box<Expression>,
     },
+    Branch(String),
     BranchIf {
         condition: Box<Expression>,
         label: String,
@@ -101,6 +102,10 @@ pub enum Expr {
         left: Box<Expression>,
         right: Box<Expression>,
     },
+    Assign {
+        name: String,
+        value: Box<Expression>,
+    },
     LocalTee {
         name: String,
         value: Box<Expression>,
@@ -123,6 +128,9 @@ pub enum Expr {
         if_true: Box<Expression>,
         if_false: Option<Box<Expression>>
     },
+    Return {
+        value: Option<Box<Expression>>
+    },
     Error,
 }
 
@@ -157,6 +165,9 @@ pub enum BinOp {
     Ge,
     Lt,
     Le,
+    Lsl,
+    Lsr,
+    Asr,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

src/constfold.rs

@@ -83,6 +83,9 @@ fn fold_expr(expr: &mut ast::Expression) {
                         Le => (left <= right) as i32,
                         Gt => (left > right) as i32,
                         Ge => (left >= right) as i32,
+                        Lsl => left << right,
+                        Lsr => ((left as u32) >> right) as i32,
+                        Asr => left >> right
                     };
                     expr.expr = ast::Expr::I32Const(result);
                 }
@@ -93,7 +96,7 @@ fn fold_expr(expr: &mut ast::Expression) {
                         Sub => F32Const(left - right),
                         Mul => F32Const(left * right),
                         Div => F32Const(left / right),
-                        Rem | And | Or | Xor => return,
+                        Rem | And | Or | Xor | Lsl | Lsr | Asr => return,
                         Eq => I32Const((left == right) as i32),
                         Ne => I32Const((left != right) as i32),
                         Lt => I32Const((left < right) as i32),
@@ -106,8 +109,10 @@ fn fold_expr(expr: &mut ast::Expression) {
             }
         }
         ast::Expr::I32Const(_) | ast::Expr::F32Const(_) | ast::Expr::Variable { .. } => (),
+        ast::Expr::Assign { ref mut value, .. } => fold_expr(value),
         ast::Expr::LocalTee { ref mut value, .. } => fold_expr(value),
         ast::Expr::Loop { ref mut block, .. } => fold_expr(block),
+        ast::Expr::Branch(_) => (),
         ast::Expr::BranchIf {
             ref mut condition, ..
         } => fold_expr(condition),
@@ -146,6 +151,8 @@ fn fold_expr(expr: &mut ast::Expression) {
                 fold_expr(if_false);
             }
         }
+        ast::Expr::Return { value: Some(ref mut value) } => fold_expr(value),
+        ast::Expr::Return { value: None } => (),
         ast::Expr::Error => unreachable!()
     }
 }

src/emit.rs

@@ -68,6 +68,11 @@ pub fn emit(script: &ast::Script) -> Vec<u8> {
         let mut exports = ExportSection::new();
         let mut code = CodeSection::new();
 
+        let mut function_map = HashMap::new();
+        for func in script.functions.iter() {
+            function_map.insert(func.name.clone(), function_map.len() as u32);
+        }
+
         for (index, func) in script.functions.iter().enumerate() {
             let type_ = *function_types.get(&function_type_key(func)).unwrap();
             functions.function(type_ as u32);
@@ -75,7 +80,7 @@ pub fn emit(script: &ast::Script) -> Vec<u8> {
                 exports.export(&func.name, Export::Function(index as u32));
             }
 
-            code.function(&emit_function(func, &globals));
+            code.function(&emit_function(func, &globals, &function_map));
         }
 
         module.section(&functions);
@@ -109,28 +114,38 @@ fn function_type_key(func: &ast::Function) -> FunctionTypeKey {
 struct FunctionContext<'a> {
     function: &'a mut Function,
     globals: &'a HashMap<&'a str, u32>,
+    functions: &'a HashMap<String, u32>,
     locals: &'a HashMap<String, u32>,
     labels: Vec<String>,
     deferred_inits: HashMap<&'a str, &'a ast::Expression>,
 }
 
-fn emit_function(func: &ast::Function, globals: &HashMap<&str, u32>) -> Function {
+fn emit_function(
+    func: &ast::Function,
+    globals: &HashMap<&str, u32>,
+    functions: &HashMap<String, u32>,
+) -> 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 locals: HashMap<String, u32> = locals
+    let mut local_map: HashMap<String, u32> = HashMap::new();
-        .into_iter()
+
-        .enumerate()
+    for (ref name, _) in func.params.iter() {
-        .map(|(index, (name, _))| (name, index as u32))
+        local_map.insert(name.clone(), local_map.len() as u32);
-        .collect();
+    }
+
+    for (name, _) in locals {
+        local_map.insert(name, local_map.len() as u32);
+    }
 
     let mut context = FunctionContext {
         function: &mut function,
         globals,
-        locals: &locals,
+        functions,
+        locals: &local_map,
         labels: vec![],
         deferred_inits: HashMap::new(),
     };
@@ -180,7 +195,9 @@ fn collect_locals_expr<'a>(expr: &ast::Expression, locals: &mut Vec<(String, ast
             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::Cast { value, .. } => collect_locals_expr(value, locals),
@@ -210,6 +227,8 @@ fn collect_locals_expr<'a>(expr: &ast::Expression, locals: &mut Vec<(String, ast
                 collect_locals_expr(if_false, locals);
             }
         }
+        ast::Expr::Return { value: Some(value) } => collect_locals_expr(value, locals),
+        ast::Expr::Return { value: None } => (),
         ast::Expr::Error => unreachable!(),
     }
 }
@@ -318,12 +337,15 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
                 (I32, Le) => Instruction::I32LeS,
                 (I32, Gt) => Instruction::I32GtS,
                 (I32, Ge) => Instruction::I32GeS,
+                (I32, Lsl) => Instruction::I32Shl,
+                (I32, Lsr) => Instruction::I32ShrU,
+                (I32, Asr) => Instruction::I32ShrS,
 
                 (F32, Add) => Instruction::F32Add,
                 (F32, Sub) => Instruction::F32Sub,
                 (F32, Mul) => Instruction::F32Mul,
                 (F32, Div) => Instruction::F32Div,
-                (F32, Rem | And | Or | Xor) => unreachable!(),
+                (F32, Rem | And | Or | Xor | Lsl | Lsr | Asr) => unreachable!(),
                 (F32, Eq) => Instruction::F32Eq,
                 (F32, Ne) => Instruction::F32Neq,
                 (F32, Lt) => Instruction::F32Lt,
@@ -335,6 +357,17 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
                 (F64, _) => todo!(),
             });
         }
+        ast::Expr::Branch(label) => {
+            let depth = ctx
+                .labels
+                .iter()
+                .rev()
+                .enumerate()
+                .find(|(_, l)| *l == label)
+                .unwrap()
+                .0;
+            ctx.function.instruction(&Instruction::Br(depth as u32));
+        }
         ast::Expr::BranchIf {
             condition, label, ..
         } => {
@@ -355,6 +388,18 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
         ast::Expr::F32Const(v) => {
             ctx.function.instruction(&Instruction::F32Const(*v));
         }
+        ast::Expr::Assign { name, value, .. } => {
+            emit_expression(ctx, value);
+            if let Some(local_index) = ctx.locals.get(name) {
+                ctx.function
+                    .instruction(&Instruction::LocalSet(*local_index));
+            } else if let Some(global_index) = ctx.globals.get(name.as_str()) {
+                ctx.function
+                    .instruction(&Instruction::GlobalSet(*global_index));
+            } else {
+                unreachable!();
+            }
+        }
         ast::Expr::LocalTee { name, value, .. } => {
             emit_expression(ctx, value);
             let index = ctx.locals.get(name).unwrap();
@@ -379,6 +424,7 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
             } else if let Some(index) = ctx.globals.get(name.as_str()) {
                 ctx.function.instruction(&Instruction::GlobalGet(*index));
             } else {
+                dbg!(name);
                 unreachable!()
             }
         }
@@ -396,13 +442,19 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
             }
         }
         ast::Expr::FuncCall { name, params, .. } => {
-            let mut types = vec![];
             for param in params {
-                types.push(param.type_.unwrap());
                 emit_expression(ctx, param);
             }
-            ctx.function
+            if let Some(index) = ctx.functions.get(name) {
-                .instruction(&builtin_function(name, &types).unwrap());
+                ctx.function.instruction(&Instruction::Call(*index));
+            } else {
+                let mut types = vec![];
+                for param in params {
+                    types.push(param.type_.unwrap());
+                }
+                ctx.function
+                    .instruction(&builtin_function(name, &types).unwrap());
+            }
         }
         ast::Expr::Select {
             condition,
@@ -434,6 +486,13 @@ fn emit_expression<'a>(ctx: &mut FunctionContext<'a>, expr: &'a ast::Expression)
                     ctx.function.instruction(&Instruction::Drop);
                 }
             }
+            ctx.function.instruction(&Instruction::End);
+        }
+        ast::Expr::Return { value } => {
+            if let Some(value) = value {
+                emit_expression(ctx, value);
+                ctx.function.instruction(&Instruction::Return);
+            }
         }
         ast::Expr::Error => unreachable!(),
     }

src/main.rs

@@ -26,13 +26,13 @@ fn main() -> Result<()> {
 
     constfold::fold_script(&mut script);
     if let Err(_) = typecheck::tc_script(&mut script, &input) {
-        bail!("Parse failed");
+        bail!("Type check failed");
     }
     let wasm = emit::emit(&script);
 
     wasmparser::validate(&wasm)?;
 
-    filename.set_extension("uw8");
+    filename.set_extension("wasm");
     File::create(filename)?.write_all(&wasm)?;
 
     println!("Size of code section: {} bytes", code_section_size(&wasm)?);

src/parser.rs

@@ -14,12 +14,14 @@ enum Token {
     Global,
     Mut,
     Loop,
+    Branch,
     BranchIf,
     Defer,
     As,
     Select,
     If,
     Else,
+    Return,
     Ident(String),
     Str(String),
     Int(i32),
@@ -39,12 +41,14 @@ impl fmt::Display for Token {
             Token::Global => write!(f, "global"),
             Token::Mut => write!(f, "mut"),
             Token::Loop => write!(f, "loop"),
+            Token::Branch => write!(f, "branch"),
             Token::BranchIf => write!(f, "branch_if"),
             Token::Defer => write!(f, "defer"),
             Token::As => write!(f, "as"),
             Token::Select => write!(f, "select"),
             Token::If => write!(f, "if"),
             Token::Else => write!(f, "else"),
+            Token::Return => write!(f, "return"),
             Token::Ident(s) => write!(f, "{}", s),
             Token::Str(s) => write!(f, "{:?}", s),
             Token::Int(v) => write!(f, "{}", v),
@@ -191,12 +195,14 @@ fn lexer() -> impl Parser<char, Vec<(Token, Span)>, Error = Simple<char>> {
         "global" => Token::Global,
         "mut" => Token::Mut,
         "loop" => Token::Loop,
+        "branch" => Token::Branch,
         "branch_if" => Token::BranchIf,
         "defer" => Token::Defer,
         "as" => Token::As,
         "select" => Token::Select,
         "if" => Token::If,
-        "Else" => Token::Else,
+        "else" => Token::Else,
+        "return" => Token::Return,
         _ => Token::Ident(ident),
     });
 
@@ -304,6 +310,10 @@ fn block_parser() -> impl Parser<Token, ast::Expression, Error = Simple<Token>>
 
             block_expression = Some(block_expr.clone());
 
+            let branch = just(Token::Branch)
+                .ignore_then(ident)
+                .map(|label| ast::Expr::Branch(label));
+
             let branch_if = just(Token::BranchIf)
                 .ignore_then(expression.clone())
                 .then_ignore(just(Token::Ctrl(':')))
@@ -341,6 +351,16 @@ fn block_parser() -> impl Parser<Token, ast::Expression, Error = Simple<Token>>
                 })
                 .boxed();
 
+            let assign = ident
+                .clone()
+                .then_ignore(just(Token::Op("=".to_string())))
+                .then(expression.clone())
+                .map(|(name, value)| ast::Expr::Assign {
+                    name,
+                    value: Box::new(value),
+                })
+                .boxed();
+
             let select = just(Token::Select)
                 .ignore_then(
                     expression
@@ -369,15 +389,24 @@ fn block_parser() -> impl Parser<Token, ast::Expression, Error = Simple<Token>>
                 .map(|(name, params)| ast::Expr::FuncCall { name, params })
                 .boxed();
 
+            let return_ = just(Token::Return)
+                .ignore_then(expression.clone().or_not())
+                .map(|value| ast::Expr::Return {
+                    value: value.map(Box::new),
+                });
+
             let atom = val
                 .or(tee)
                 .or(function_call)
-                .or(variable)
+                .or(assign)
                 .or(local_tee)
+                .or(variable)
                 .or(block_expr)
+                .or(branch)
                 .or(branch_if)
                 .or(let_)
                 .or(select)
+                .or(return_)
                 .map_with_span(|expr, span| expr.with_span(span))
                 .or(expression
                     .clone()
@@ -540,7 +569,28 @@ fn block_parser() -> impl Parser<Token, ast::Expression, Error = Simple<Token>>
                 })
                 .boxed();
 
-            let op_cmp = op_sum
+            let op_shift = op_sum
+                .clone()
+                .then(
+                    just(Token::Op("<<".to_string()))
+                        .to(ast::BinOp::Lsl)
+                        .or(just(Token::Op(">>".to_string())).to(ast::BinOp::Lsr))
+                        .or(just(Token::Op(">>>".to_string())).to(ast::BinOp::Asr))
+                        .then(op_sum.clone())
+                        .repeated(),
+                )
+                .foldl(|left, (op, right)| {
+                    let span = left.span.start..right.span.end;
+                    ast::Expr::BinOp {
+                        op,
+                        left: Box::new(left),
+                        right: Box::new(right),
+                    }
+                    .with_span(span)
+                })
+                .boxed();
+
+            let op_cmp = op_shift
                 .clone()
                 .then(
                     just(Token::Op("==".to_string()))
@@ -550,7 +600,7 @@ fn block_parser() -> impl Parser<Token, ast::Expression, Error = Simple<Token>>
                         .or(just(Token::Op("<=".to_string())).to(ast::BinOp::Le))
                         .or(just(Token::Op(">".to_string())).to(ast::BinOp::Gt))
                         .or(just(Token::Op(">=".to_string())).to(ast::BinOp::Ge))
-                        .then(op_sum.clone())
+                        .then(op_shift.clone())
                         .repeated(),
                 )
                 .foldl(|left, (op, right)| {

src/typecheck.rs

@@ -17,8 +17,10 @@ pub fn tc_script(script: &mut ast::Script, source: &str) -> Result<()> {
     let mut context = Context {
         source,
         global_vars: HashMap::new(),
+        functions: HashMap::new(),
         local_vars: HashMap::new(),
         block_stack: Vec::new(),
+        return_type: None,
     };
 
     let mut result = Ok(());
@@ -64,6 +66,23 @@ pub fn tc_script(script: &mut ast::Script, source: &str) -> Result<()> {
         }
     }
 
+    for f in &script.functions {
+        let params = f.params.iter().map(|(_, t)| *t).collect();
+        if let Some(fnc) = context.functions.get(&f.name) {
+            result =
+                report_duplicate_definition("Function already defined", &f.span, &fnc.span, source);
+        } else {
+            context.functions.insert(
+                f.name.clone(),
+                FunctionType {
+                    params,
+                    type_: f.type_,
+                    span: f.span.clone(),
+                },
+            );
+        }
+    }
+
     for f in &mut script.functions {
         context.local_vars.clear();
         for (name, type_) in &f.params {
@@ -84,18 +103,31 @@ pub fn tc_script(script: &mut ast::Script, source: &str) -> Result<()> {
                 );
             }
         }
+        context.return_type = f.type_;
 
         tc_expression(&mut context, &mut f.body)?;
+
+        if f.body.type_ != f.type_ {
+            result = type_mismatch(f.type_, &f.span, f.body.type_, &f.body.span, source);
+        }
     }
 
     result
 }
 
+struct FunctionType {
+    span: Span,
+    params: Vec<ast::Type>,
+    type_: Option<ast::Type>,
+}
+
 struct Context<'a> {
     source: &'a str,
     global_vars: Vars,
+    functions: HashMap<String, FunctionType>,
     local_vars: Vars,
     block_stack: Vec<String>,
+    return_type: Option<ast::Type>,
 }
 
 fn report_duplicate_definition(
@@ -123,7 +155,7 @@ fn report_duplicate_definition(
 }
 
 fn type_mismatch(
-    type1: ast::Type,
+    type1: Option<ast::Type>,
     span1: &Span,
     type2: Option<ast::Type>,
     span2: &Span,
@@ -133,7 +165,12 @@ fn type_mismatch(
         .with_message("Type mismatch")
         .with_label(
             Label::new(span1.clone())
-                .with_message(format!("Expected type {:?}...", type1))
+                .with_message(format!(
+                    "Expected type {:?}...",
+                    type1
+                        .map(|t| format!("{:?}", t))
+                        .unwrap_or("void".to_string())
+                ))
                 .with_color(Color::Yellow),
         )
         .with_label(
@@ -180,11 +217,25 @@ fn unknown_variable(span: &Span, source: &str) -> Result<()> {
     Err(())
 }
 
+fn missing_label(span: &Span, source: &str) -> Result<()> {
+    Report::build(ReportKind::Error, (), span.start)
+        .with_message("Label not found")
+        .with_label(
+            Label::new(span.clone())
+                .with_message("Label not found")
+                .with_color(Color::Red),
+        )
+        .finish()
+        .eprint(Source::from(source))
+        .unwrap();
+    return Err(());
+}
+
 fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()> {
     expr.type_ = match expr.expr {
         ast::Expr::Block {
             ref mut statements,
-            ref mut final_expression
+            ref mut final_expression,
         } => {
             for stmt in statements {
                 tc_expression(context, stmt)?;
@@ -207,7 +258,7 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
                 if let Some(type_) = type_ {
                     if Some(*type_) != value.type_ {
                         return type_mismatch(
-                            *type_,
+                            Some(*type_),
                             &expr.span,
                             value.type_,
                             &value.span,
@@ -266,7 +317,13 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
             tc_mem_location(context, mem_location)?;
             tc_expression(context, value)?;
             if value.type_ != Some(I32) {
-                return type_mismatch(I32, &expr.span, value.type_, &value.span, context.source);
+                return type_mismatch(
+                    Some(I32),
+                    &expr.span,
+                    value.type_,
+                    &value.span,
+                    context.source,
+                );
             }
             None
         }
@@ -289,7 +346,7 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
             if let Some(type_) = left.type_ {
                 if left.type_ != right.type_ {
                     return type_mismatch(
-                        type_,
+                        Some(type_),
                         &left.span,
                         right.type_,
                         &right.span,
@@ -302,10 +359,10 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
             use ast::BinOp::*;
             match op {
                 Add | Sub | Mul | Div => left.type_,
-                Rem | And | Or | Xor => {
+                Rem | And | Or | Xor | Lsl | Lsr | Asr => {
                     if left.type_ != Some(I32) {
                         return type_mismatch(
-                            I32,
+                            Some(I32),
                             &left.span,
                             left.type_,
                             &left.span,
@@ -329,6 +386,32 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
                 return unknown_variable(&expr.span, context.source);
             }
         }
+        ast::Expr::Assign {
+            ref name,
+            ref mut value,
+        } => {
+            tc_expression(context, value)?;
+            if let Some(&Var {
+                type_, ref span, ..
+            }) = context
+                .local_vars
+                .get(name)
+                .or_else(|| context.global_vars.get(name))
+            {
+                if value.type_ != Some(type_) {
+                    return type_mismatch(
+                        Some(type_),
+                        span,
+                        value.type_,
+                        &value.span,
+                        context.source,
+                    );
+                }
+            } else {
+                return unknown_variable(&expr.span, context.source);
+            }
+            None
+        }
         ast::Expr::LocalTee {
             ref name,
             ref mut value,
@@ -339,7 +422,13 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
             }) = context.local_vars.get(name)
             {
                 if value.type_ != Some(type_) {
-                    return type_mismatch(type_, span, value.type_, &value.span, context.source);
+                    return type_mismatch(
+                        Some(type_),
+                        span,
+                        value.type_,
+                        &value.span,
+                        context.source,
+                    );
                 }
                 Some(type_)
             } else {
@@ -355,6 +444,12 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
             context.block_stack.pop();
             block.type_
         }
+        ast::Expr::Branch(ref label) => {
+            if !context.block_stack.contains(label) {
+                return missing_label(&expr.span, context.source);
+            }
+            None
+        }
         ast::Expr::BranchIf {
             ref mut condition,
             ref label,
@@ -362,7 +457,7 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
             tc_expression(context, condition)?;
             if condition.type_ != Some(I32) {
                 return type_mismatch(
-                    I32,
+                    Some(I32),
                     &expr.span,
                     condition.type_,
                     &condition.span,
@@ -370,17 +465,7 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
                 );
             }
             if !context.block_stack.contains(label) {
-                Report::build(ReportKind::Error, (), expr.span.start)
+                return missing_label(&expr.span, context.source);
-                    .with_message("Label not found")
-                    .with_label(
-                        Label::new(expr.span.clone())
-                            .with_message("Label not found")
-                            .with_color(Color::Red),
-                    )
-                    .finish()
-                    .eprint(Source::from(context.source))
-                    .unwrap();
-                return Err(());
             }
             None
         }
@@ -398,7 +483,15 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
             ref name,
             ref mut params,
         } => {
-            if let Some((ptypes, rtype)) = builtin_function_types(name) {
+            for param in params.iter_mut() {
+                tc_expression(context, param)?;
+            }
+            if let Some((ptypes, rtype)) = context
+                .functions
+                .get(name)
+                .map(|fnc| (fnc.params.as_slice(), fnc.type_))
+                .or_else(|| builtin_function_types(name))
+            {
                 if params.len() != ptypes.len() {
                     Report::build(ReportKind::Error, (), expr.span.start)
                         .with_message(format!(
@@ -420,11 +513,10 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
                         .unwrap();
                     return Err(());
                 }
-                for (ptype, param) in ptypes.iter().zip(params.iter_mut()) {
+                for (ptype, param) in ptypes.iter().zip(params.iter()) {
-                    tc_expression(context, param)?;
                     if param.type_ != Some(*ptype) {
                         return type_mismatch(
-                            *ptype,
+                            Some(*ptype),
                             &expr.span,
                             param.type_,
                             &param.span,
@@ -457,17 +549,17 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
             tc_expression(context, if_false)?;
             if condition.type_ != Some(ast::Type::I32) {
                 return type_mismatch(
-                    I32,
+                    Some(I32),
                     &condition.span,
                     condition.type_,
                     &condition.span,
                     context.source,
                 );
             }
-            if let Some(true_type) = if_true.type_ {
+            if if_true.type_.is_some() {
                 if if_true.type_ != if_false.type_ {
                     return type_mismatch(
-                        true_type,
+                        if_true.type_,
                         &if_true.span,
                         if_false.type_,
                         &if_false.span,
@@ -482,15 +574,20 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
         ast::Expr::If {
             ref mut condition,
             ref mut if_true,
-            ref mut if_false
+            ref mut if_false,
         } => {
             tc_expression(context, condition)?;
             tc_expression(context, if_true)?;
             if let Some(ref mut if_false) = if_false {
                 tc_expression(context, if_false)?;
                 if if_true.type_ != if_false.type_ {
-                    // TODO: report type mismatch?
+                    return type_mismatch(
-                    None
+                        if_true.type_,
+                        &if_true.span,
+                        if_false.type_,
+                        &if_false.span,
+                        context.source,
+                    );
                 } else {
                     if_true.type_
                 }
@@ -498,6 +595,21 @@ fn tc_expression(context: &mut Context, expr: &mut ast::Expression) -> Result<()
                 None
             }
         }
+        ast::Expr::Return { ref mut value } => {
+            if let Some(ref mut value) = value {
+                tc_expression(context, value)?;
+                if value.type_ != context.return_type {
+                    return type_mismatch(
+                        context.return_type,
+                        &expr.span,
+                        value.type_,
+                        &value.span,
+                        context.source,
+                    );
+                }
+            }
+            None
+        }
         ast::Expr::Error => unreachable!(),
     };
     Ok(())
@@ -511,7 +623,7 @@ fn tc_mem_location<'a>(
     tc_expression(context, &mut mem_location.right)?;
     if mem_location.left.type_ != Some(I32) {
         return type_mismatch(
-            I32,
+            Some(I32),
             &mem_location.left.span,
             mem_location.left.type_,
             &mem_location.left.span,

uw8loader.hw

@@ -0,0 +1,64 @@
+import "uw8.ram" memory(8);
+
+export fn load_uw8(module_start: i32, module_end: i32, base_start: i32, base_end: i32) -> i32 {
+    if ?module_start == 0 {
+        let defer length = module_end - module_start;
+        copy(base_end, module_start, length);
+        return length;
+    }
+
+    copy(base_end, base_start, 8);
+    base_start = base_start + 8;
+    let dest = base_end + 8;
+    let src = module_start + 1;
+
+    loop sections {
+        if src < module_end & (base_start >= base_end | ?src <= ?base_start) {
+            let length2 = copy_section(dest, src);
+            dest = dest + length2;
+            if base_start < base_end & ?src == ?base_start {
+                base_start = base_start + section_size(base_start);
+            }
+            src = src + length2;
+            branch sections;
+        }
+
+        if base_start < base_end {
+            let length3 = copy_section(dest, base_start);
+            dest = dest + length3;
+            base_start = base_start + length3;
+            branch sections;
+        }
+    }
+
+    dest
+}
+
+fn section_size(ptr: i32) -> i32 {
+    let p = ptr + 1;
+    let l = 0;
+    let shift = 0;
+    loop size {
+        let b = ?p;
+        p = p + 1;
+        l = l | ((b & 127) << shift);
+        shift = shift + 7;
+        branch_if b & 128: size;
+    }
+    p - ptr + l
+}
+
+fn copy_section(dest: i32, src: i32) -> i32 {
+    let defer length = section_size(src);
+    copy(dest, src, length);
+    length
+}
+
+fn copy(dest: i32, src: i32, len: i32) {
+    if len > 0 {
+        loop bytes {
+            ?(dest + (len := len - 1)) = ?(src + len);
+            branch_if len: bytes
+        }
+    }
+}