cc: implement proper support for return statements

atirut-w · Jul 27, 2024 · b4de0bd · b4de0bd
1 parent a5adfc3
commit b4de0bd
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 10 deletions.
diff --git a/src/cc/analyzer.cpp b/src/cc/analyzer.cpp
@@ -1,5 +1,6 @@
 #include "CParser.h"
 #include <analyzer.hpp>
+#include <any>
 
 using namespace std;
 using namespace antlr4;
@@ -20,7 +21,27 @@ any Analyzer::visitFunctionDefinition(CParser::FunctionDefinitionContext *ctx)
     }
     current_function = &meta.functions[name];
 
-    return visitChildren(ctx);
+    if (auto *itemlist_ctx = ctx->compoundStatement()->blockItemList())
+    {
+        for (auto *item_ctx : itemlist_ctx->blockItem())
+        {
+            // We check for return statements here because we only want to check if this function *ends* with a return statement
+            if (auto *statement_ctx = item_ctx->statement())
+            {
+                if (auto *jump_statement_ctx = statement_ctx->jumpStatement())
+                {
+                    if (jump_statement_ctx->Return())
+                    {
+                        current_function->has_return = true;
+                    }
+                }
+            }
+
+            visit(item_ctx);
+        }
+    }
+
+    return any();
 }
 
 any Analyzer::visitDeclaration(CParser::DeclarationContext *ctx)

diff --git a/src/cc/codegen.cpp b/src/cc/codegen.cpp
@@ -55,6 +55,20 @@ CodeGen::CodeGen(ProgramMeta &program_meta, std::ostream &output) : program_meta
     output << "\t; Generated by ZDK C Compiler\n\n";
 }
 
+void CodeGen::teardown_frame()
+{
+    // TODO: clause for arguments
+    if (current_function->local_alloc > 0 || false)
+    {
+        output << "\tld sp, ix\n";
+        output << "\tpop ix\n";
+        if (current_function->local_alloc > 0)
+        {
+            output << "\tpop iy\n";
+        }
+    }
+}
+
 any CodeGen::visitFunctionDefinition(CParser::FunctionDefinitionContext *ctx)
 {
     string name = ctx->declarator()->directDeclarator()->directDeclarator()->Identifier()->getText();
@@ -91,18 +105,27 @@ any CodeGen::visitFunctionDefinition(CParser::FunctionDefinitionContext *ctx)
         }
     }
 
-    // TODO: clause for arguments
-    if (current_function->local_alloc > 0 || false)
+    if (!current_function->has_return)
     {
-        output << "\tld sp, ix\n";
-        output << "\tpop ix\n";
-        if (current_function->local_alloc > 0)
-        {
-            output << "\tpop iy\n";
-        }
+        teardown_frame();
+        output << "\tret\n";
     }
 
-    output << "\tret\n";
+    return any();
+}
+
+any CodeGen::visitJumpStatement(CParser::JumpStatementContext *ctx)
+{
+    if (ctx->Return())
+    {
+        visit(ctx->expression());
+        teardown_frame();
+        output << "\tret\n";
+    }
+    else
+    {
+        throw runtime_error("unsupported jump statement");
+    }
 
     return any();
 }

diff --git a/src/cc/include/analyzer.hpp b/src/cc/include/analyzer.hpp
@@ -20,6 +20,7 @@ struct FunctionMeta
     // Offset into the local frame for local variables
     std::map<std::string, LocalMeta> variables;
     int local_alloc = 0;
+    bool has_return = false; // TODO: Come up with a better name for this
 };
 
 struct ProgramMeta

diff --git a/src/cc/include/codegen.hpp b/src/cc/include/codegen.hpp
@@ -20,7 +20,10 @@ class CodeGen : public CBaseVisitor
     FunctionMeta *current_function = nullptr;
     std::ostream &output;
 
+    void teardown_frame();
+
     virtual std::any visitFunctionDefinition(CParser::FunctionDefinitionContext *ctx) override;
+    virtual std::any visitJumpStatement(CParser::JumpStatementContext *ctx) override;
 
     // All of these are for math expressions. Yes, all of them.
     virtual std::any visitPrimaryExpression(CParser::PrimaryExpressionContext *ctx) override; // Your pain starts here