added method calling and for loops

1 files changed, 299 insertions, 57 deletions

diff --git a/Compiler/CCompilerVisitor.cs b/Compiler/CCompilerVisitor.cs
index fb7df0f..abe831d 100644
--- a/Compiler/CCompilerVisitor.cs
+++ b/Compiler/CCompilerVisitor.cs
@@ -1,30 +1,33 @@
 namespace Q1.Compiler;
 
+using System.Runtime.CompilerServices;
 using Antlr4.Runtime.Tree;
 using Emulator;
 
 public class CCompilerVisitor : CGrammarBaseVisitor<string?>
 {
-    private int                            _branchCount   = 1;
-    private List<string>                   _instructions  = new();
-    private Stack<Dictionary<string, int>> _variableStack = new();
-    private int                            _memPointer    = Q1Layout.FreeMemoryStart;
+    private int                                          _branchCount   = 1;
+    private List<string>                                 _instructions  = new();
+    private Stack<Dictionary<string, VariableReference>> _variableStack = new();
+    private List<FunctionReference>                      _methods       = new();
+    private static int                                          _memLowest     = Q1Layout.FreeMemoryStart;
+    private static int                                          _memPointer    = CCompilerVisitor._memLowest;
 
     private const int ColumnWidth = 30;
 
     public readonly CommentCompilationMode CommentMode;
-    
+
     public CCompilerVisitor(CommentCompilationMode commentMode)
     {
         this.CommentMode = commentMode;
     }
 
-    private int FindVariable(string varName)
+    private VariableReference FindVariable(string varName)
     {
-        foreach (Dictionary<string, int> variables in this._variableStack)
+        foreach (Dictionary<string, VariableReference> variables in this._variableStack)
         {
-            if (variables.TryGetValue(varName, out int address))
-                return address;
+            if (variables.TryGetValue(varName, out VariableReference? reference))
+                return reference;
         }
 
         throw new CompilerException($"Variable {varName} could not be found in current scope.");
@@ -32,8 +35,8 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
 
     private int Alloc(int size)
     {
-        int loc = this._memPointer;
-        this._memPointer += size;
+        int loc = CCompilerVisitor._memPointer;
+        CCompilerVisitor._memPointer += size;
 
         return loc;
     }
@@ -52,7 +55,7 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
     {
         this._instructions.Add("");
     }
-    private void Label(string name, string? generatedComment = null)
+    private void Branch(string name, string? generatedComment = null)
     {
         this._instructions.Add($"{name}:");
         if (generatedComment != null)
@@ -63,7 +66,7 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         if (this.CommentMode.HasFlag(CommentCompilationMode.UserDefined))
             this._instructions.Add($"    ; {comment}");
     }
-    
+
     private void GeneratedComment(string comment)
     {
         if (this.CommentMode.HasFlag(CommentCompilationMode.Generated))
@@ -87,25 +90,139 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
 
     public IEnumerable<string> Compile(CGrammarParser.ProgramContext program)
     {
+        Instruction("jmp _main");
         this.Visit(program);
         return this._instructions;
     }
 
-    public override string? VisitFunction(CGrammarParser.FunctionContext context)
+    private FunctionReference CreateOrGetFunctionReference(string name, string type)
+    {
+        FunctionReference? function = this._methods.FirstOrDefault(x => x.Name == name);
+        if (function != null)
+            return function;
+        
+        function = new FunctionReference(name, type);
+        this._methods.Add(function);
+
+        return function;
+    }
+    
+    private FunctionReference GetFunctionReference(string name)
+    {
+        FunctionReference? function = this._methods.FirstOrDefault(x => x.Name == name);
+        if (function != null)
+            return function;
+
+        throw new CompilerException($"Function '{name}' has not been defined or prototyped.");
+    }
+
+    public override string? VisitFunctionPrototype(CGrammarParser.FunctionPrototypeContext context)
     {
         string functionName = context.name.Text;
-        this._variableStack.Push(new());
+        string type = "int";
+
+        List<string> parameterTypes = context._types.Select(x => x.Text).ToList();
+        List<string> parameterNames = context._params.Select(x => x.Text).ToList();
 
-        this.Label($"_{functionName}");
+        FunctionReference function = this.CreateOrGetFunctionReference(functionName, type);
+        FunctionOverload overload = function.CreateOverloadWithoutBody(parameterNames, parameterTypes);
+        
+        function.Overloads.Add(overload);
+
+        return null;
+    }
+    public override string? VisitFunctionDefinition(CGrammarParser.FunctionDefinitionContext context)
+    {
+        string functionName = context.name.Text;
+        string type = "int";
+        
+        List<string> paramTypes = context._types.Select(x => x.Text).ToList();
+        List<string> paramNames = context._params.Select(x => x.Text).ToList();
+
+        FunctionReference function = this.CreateOrGetFunctionReference(functionName, type);
+        FunctionOverload overload = function.CreateOverloadWithBody(paramNames, paramTypes);
+
+        this.Branch(function.GetBranchName(overload), function.GetBranchFriendlyName(overload));
+
+        if (paramNames.Count > 0)
+            this.Instruction("pop V0", "preserve function return address");
+        
+        Dictionary<string, VariableReference> scope = this.EnterScope();
+        foreach ((string paramName, string paramType) in paramNames.Zip(paramTypes))
+        {
+            int paramSize = 2;
+            int addr = this.Alloc(paramSize);
+            scope[paramName] = new(addr, paramSize);
+            this.Instruction($"pop [${addr:X4}]", $"param {paramType} {paramName}");
+        }
+        
+        if (paramNames.Count > 0)
+            this.Instruction("push V0", "restore function return address");
+        
         this.VisitChildren(context);
+        this.ExitScope();
+        
         this.Instruction("mov 0, V0", "fallback return 0");
         this.Instruction("ret");
         this.Blank();
 
-        this._variableStack.Pop();
+        return null;
+    }
+
+    public override string? VisitCallExpression(CGrammarParser.CallExpressionContext context)
+    {
+        string functionName = context.ID().GetText();
+
+        FunctionReference function = this.GetFunctionReference(functionName);
+        FunctionOverload overload = function.GetOverload(context._params.Count);
+
+        // We reverse arugments so when we perform pop, its in the right order
+        CGrammarParser.ExpressionContext[] argumentsReversed = context._params.Reverse().ToArray();
+        
+        for (int i = 0; i < argumentsReversed.Length; i++)
+        {
+            CGrammarParser.ExpressionContext argument = argumentsReversed[i];
+            string result = this.Visit(argument)
+                            ?? throw new CompilerException("Function call argument was not an expression");
+
+            this.Instruction($"push {result}", $"arg {overload.ParameterTypes[i]} {overload.ParameterNames[i]}");
+        }
+        
+        this.Instruction($"call {function.GetBranchName(overload)}", function.GetBranchFriendlyName(overload));
+
+        return "V0";
+    }
+
+    public override string? VisitBlock(CGrammarParser.BlockContext context)
+    {
+        this.EnterScope();
+
+        this.VisitChildren(context);
+
+        this.ExitScope();
 
         return null;
     }
+    private Dictionary<string, VariableReference> EnterScope()
+    {
+        Dictionary<string, VariableReference> scope = new();
+        this._variableStack.Push(scope);
+
+        return scope;
+    }
+    private void ExitScope()
+    {
+        foreach ((string key, VariableReference value) in this._variableStack.Peek())
+        {
+            CCompilerVisitor._memPointer -= value.Size;
+            if (CCompilerVisitor._memPointer < CCompilerVisitor._memLowest)
+                throw new CompilerException("Memory pointer is lower than mem lowest!");
+        }
+
+        // Delete scope
+        this._variableStack.Pop();
+    }
+
     public override string? VisitParenthesizedExpression(CGrammarParser.ParenthesizedExpressionContext context)
     {
         return this.Visit(context.expression());
@@ -169,8 +286,8 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         if (right == null)
             throw new CompilerException($"Right operand for {operationFriendlyName} was not an expression.");
 
-        this.Instruction($"mov {right}, V1", "right operand");
-        this.Instruction($"pop V0", "get left operand back");
+        this.Instruction($"mov {right}, V1",    "right operand");
+        this.Instruction($"pop V0",             "get left operand back");
         this.Instruction($"{operation} V0, V1", $"compute {operationFriendlyName}");
     }
     private string BaseBinaryCompoundOperation(CGrammarParser.TargetContext leftTarget, CGrammarParser.ExpressionContext rightExpr, string operation, string operationFriendlyName, string operationOutput)
@@ -184,7 +301,7 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
             throw new CompilerException($"Right operand for {operationFriendlyName} was not an expression.");
 
         // No need for stack maniupulation because left is a target and is assumed to stay constant
-        this.Instruction($"{operation} {left}, {right}", $"compute {operationFriendlyName}");
+        this.Instruction($"{operation} {left}, {right}",   $"compute {operationFriendlyName}");
         this.Instruction($"mov {operationOutput}, {left}", "store back into left param");
         return left;
     }
@@ -253,7 +370,7 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         this.Instruction($"jmp {clauseBranch}"); // only if bz is not 0
         this.Instruction("mov 1, V0", "first operand succeeded, store 1 and skip second operand");
         this.Instruction($"jmp {endBranch}");
-        this.Label(clauseBranch);
+        this.Branch(clauseBranch);
 
         string? right = this.Visit(context.right);
         if (right == null)
@@ -262,7 +379,7 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         this.Instruction($"cmp {right}, 0", "compare result from second operand to 0");
         this.Instruction("not");
         this.Instruction("mov LX, V0", "store second operand result");
-        this.Label(endBranch);
+        this.Branch(endBranch);
 
         return "V0";
     }
@@ -281,7 +398,7 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         this.Instruction($"jmp {clauseBranch}"); // only if bz is not 0
         this.Instruction("mov 0, V0", "first operand failed, store 0 and skip second");
         this.Instruction($"jmp {endBranch}");
-        this.Label(clauseBranch);
+        this.Branch(clauseBranch);
 
         string? right = this.Visit(context.right);
         if (right == null)
@@ -290,7 +407,7 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         this.Instruction($"cmp {right}, 0", "compare result from second operand to 0");
         this.Instruction("not");
         this.Instruction("mov LX, V0", "store second operand result");
-        this.Label(endBranch);
+        this.Branch(endBranch);
 
         return "V0";
     }
@@ -303,7 +420,7 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         this.Instruction("bz");
         this.Instruction($"jmp {endBranch}");
         this.Instruction($"gt V0, V1");
-        this.Label(endBranch);
+        this.Branch(endBranch);
 
         return "LX";
     }
@@ -315,7 +432,7 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         this.Instruction("bz");
         this.Instruction($"jmp {endBranch}");
         this.Instruction($"lt V0, V1");
-        this.Label(endBranch);
+        this.Branch(endBranch);
 
         return "LX";
     }
@@ -347,18 +464,19 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
 
     public override string? VisitVariableDeclaration(CGrammarParser.VariableDeclarationContext context)
     {
-        int addr = this.Alloc(2);
+        int size = 2;
+        int addr = this.Alloc(size);
         string name = context.name.Text;
 
-        Dictionary<string, int> scope = this._variableStack.Peek();
+        Dictionary<string, VariableReference> scope = this._variableStack.Peek();
         if (scope.ContainsKey(name))
             throw new CompilerException($"Variable {name} already exists in current scope.");
 
-        scope[name] = addr;
+        scope[name] = new VariableReference(addr, size);
         if (context.value != null)
         {
-            string? initialValue = this.Visit(context.value)
-                                   ?? throw new CompilerException("Initial variable definition value was not an expression");
+            string initialValue = this.Visit(context.value)
+                                  ?? throw new CompilerException("Initial variable definition value was not an expression");
             this.Instruction($"mov {initialValue}, [${addr:X4}]", $"int {name} = {initialValue}");
         }
         else
@@ -369,24 +487,24 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
     public override string? VisitVariableExpression(CGrammarParser.VariableExpressionContext context)
     {
         string varName = context.ID().GetText();
-        int addr = this.FindVariable(varName);
+        VariableReference reference = this.FindVariable(varName);
 
-        return $"[${addr:X4}]";
+        return $"[${reference.Address:X4}]";
     }
 
     public override string? VisitVariableTarget(CGrammarParser.VariableTargetContext context)
     {
         string varName = context.ID().GetText();
-        int addr = this.FindVariable(varName);
+        VariableReference reference = this.FindVariable(varName);
 
-        return $"[${addr:X4}]";
+        return $"[${reference.Address:X4}]";
     }
     public override string? VisitAssignmentExpression(CGrammarParser.AssignmentExpressionContext context)
     {
         string? target = this.Visit(context.target());
         string? value = this.Visit(context.expression());
 
-        this.Instruction($"mov {value}, {target}");
+        this.Instruction($"mov {value}, {target}", $"{target} = {value}");
         return target;
     }
 
@@ -505,12 +623,12 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         this.Instruction($"pop V0");
         return "V0";
     }
-
-    public override string? VisitCommaExpression(CGrammarParser.CommaExpressionContext context)
-    {
-        this.Visit(context.left);
-        return this.Visit(context.right);
-    }
+    //
+    // public override string? VisitCommaExpression(CGrammarParser.CommaExpressionContext context)
+    // {
+    //     this.Visit(context.left);
+    //     return this.Visit(context.right);
+    // }
 
     public override string? VisitTernaryExpression(CGrammarParser.TernaryExpressionContext context)
     {
@@ -526,18 +644,18 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
 
         this.Instruction($"bz {cond}", "check ternary condition");
         this.Instruction($"jmp {thenBranch}");
-        
+
         result = this.Visit(context.@else)
                  ?? throw new CompilerException("Else branch was not an expression");
         this.Instruction($"mov {result}, V0", "store ternary fail result");
         this.Instruction($"jmp {endBranch}");
 
-        this.Label(thenBranch);
+        this.Branch(thenBranch);
         result = this.Visit(context.then)
                  ?? throw new CompilerException("Else branch was not an expression");
         this.Instruction($"mov {result}, V0", "store ternary success result");
 
-        this.Label(endBranch);
+        this.Branch(endBranch);
 
         return "V0";
     }
@@ -547,24 +665,24 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         string cond = this.Visit(context.expression())
                       ?? throw new CompilerException("If condition was not an expression");
         string endBranch = this.GenerateBranchLabel("end");
-        
+
         this.Instruction($"bz {cond}", "check if condition");
-        
+
         if (context.@else != null)
         {
             // Else branch is first since it saves us from having to perform a NOT operation.
             // If cond, jump to then, else fall through to else, then jump to end
 
             string thenBranch = this.GenerateBranchLabel("then");
-            
+
             this.Instruction($"jmp {thenBranch}");
 
             this.GeneratedComment("if failed");
             if (context.@else != null)
                 this.Visit(context.@else);
             this.Instruction($"jmp {endBranch}");
-        
-            this.Label(thenBranch, "if succeeded");
+
+            this.Branch(thenBranch, "if succeeded");
             this.Visit(context.then);
         }
         else
@@ -573,9 +691,127 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
             this.Instruction($"jmp {endBranch}", "if succeeds, continue, else jump to end");
             this.Visit(context.then);
         }
-        
-        this.Label(endBranch);
-        
+
+        this.Branch(endBranch);
+
+        return null;
+    }
+
+    public override string? VisitForExprStatement(CGrammarParser.ForExprStatementContext context)
+    {
+        string bodyBranch = this.GenerateBranchLabel("body");
+        string endBranch = this.GenerateBranchLabel("end");
+
+        context.LoopBodyBranch = bodyBranch;
+        context.LoopEndBranch  = endBranch;
+
+        this.EnterScope();
+        this.Visit(context.initializer);
+
+        this.Branch(bodyBranch);
+        string cond = this.Visit(context.cond)
+                      ?? throw new CompilerException("For condition was not an expression");
+
+        this.Instruction($"not {cond}");
+        this.Instruction($"bz");
+        this.Instruction($"jmp {endBranch}", "if for condition fails, jump to end");
+
+        this.Visit(context.body);
+        this.Visit(context.post);
+
+        this.Branch(endBranch);
+        this.ExitScope();
+
+        return null;
+    }
+    public override string? VisitForDeclStatement(CGrammarParser.ForDeclStatementContext context)
+    {
+        string bodyBranch = this.GenerateBranchLabel("body");
+        string endBranch = this.GenerateBranchLabel("end");
+
+        context.LoopBodyBranch = bodyBranch;
+        context.LoopEndBranch  = endBranch;
+
+        this.EnterScope();
+        this.Visit(context.initializer);
+
+        this.Branch(bodyBranch);
+        string cond = this.Visit(context.cond)
+                      ?? throw new CompilerException("For condition was not an expression");
+
+        this.Instruction($"not {cond}");
+        this.Instruction($"bz");
+        this.Instruction($"jmp {endBranch}", "if for condition fails, jump to end");
+
+        this.Visit(context.body);
+        this.Visit(context.post);
+
+        this.Branch(endBranch);
+        this.ExitScope();
+
+        return null;
+    }
+    public override string? VisitWhileStatement(CGrammarParser.WhileStatementContext context)
+    {
+        string bodyBranch = this.GenerateBranchLabel("body");
+        string endBranch = this.GenerateBranchLabel("end");
+
+        context.LoopBodyBranch = bodyBranch;
+        context.LoopEndBranch  = endBranch;
+
+        this.Branch(bodyBranch);
+        string cond = this.Visit(context.cond)
+                      ?? throw new CompilerException("While condition was not an expression");
+
+        this.Instruction($"not {cond}");
+        this.Instruction($"bz");
+        this.Instruction($"jmp {endBranch}", "if while condition fails, jump to end");
+
+        this.Visit(context.body);
+
+        this.Branch(endBranch);
+
+        return null;
+    }
+
+    public override string? VisitDoWhileStatement(CGrammarParser.DoWhileStatementContext context)
+    {
+        string bodyBranch = this.GenerateBranchLabel("body");
+        string endBranch = this.GenerateBranchLabel("end");
+
+        context.LoopBodyBranch = bodyBranch;
+        context.LoopEndBranch  = endBranch;
+
+        this.Branch(bodyBranch);
+        this.Visit(context.body);
+
+        string cond = this.Visit(context.cond)
+                      ?? throw new CompilerException("While condition was not an expression");
+
+        this.Instruction($"bz {cond}");
+        this.Instruction($"jmp {bodyBranch}", "if while condition succeeds, jump back to start");
+
+        this.Branch(endBranch);
+
+        return null;
+    }
+
+    public override string? VisitBreakStatement(CGrammarParser.BreakStatementContext context)
+    {
+        if (context.LoopEndBranch == null)
+            throw new CompilerException("Cannot break, not inside loop with end");
+
+        this.Instruction($"jmp {context.LoopEndBranch}", "break");
+
+        return null;
+    }
+    public override string? VisitContinueStatement(CGrammarParser.ContinueStatementContext context)
+    {
+        if (context.LoopBodyBranch == null)
+            throw new CompilerException("Cannot skip body (continue), not inside loop with body");
+
+        this.Instruction($"jmp {context.LoopBodyBranch}", "continue");
+
         return null;
     }
 
@@ -584,17 +820,23 @@ public class CCompilerVisitor : CGrammarBaseVisitor<string?>
         if (node.Symbol.Type == CGrammarParser.COMMENT)
         {
             // strip // and \n and whitespace
-            string comment = node.Symbol.Text
-                .TrimStart("//").TrimEnd('\n').Trim();
-            
+            string comment = node
+                .Symbol.Text
+                .TrimStart("//")
+                .TrimEnd('\n')
+                .Trim();
+
             this.UserComment(comment);
         }
-        
+
         return base.VisitErrorNode(node);
     }
 
     public override string? Visit(IParseTree tree)
     {
+        if (tree is Context childCtx && tree.Parent is Context ctx)
+            ctx.SyncProps(childCtx);
+
         return base.Visit(tree);
     }