feat(memory): add tracked MALLOC and FREE builtins

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

4 files changed, 155 insertions, 0 deletions

diff --git a/lib/CodeGen/Builtins.h b/lib/CodeGen/Builtins.h
index e7f60e4..69361cf 100644
--- a/lib/CodeGen/Builtins.h
+++ b/lib/CodeGen/Builtins.h
@@ -26,6 +26,8 @@ inline constexpr BuiltinInfo Builtins[] = {
     {"NUM_TO_STR", 1, "STRING"},
     {"STR_TO_NUM", 1, "REAL"},
     {"EOF",        1, "BOOLEAN"},
+    {"MALLOC",     1, "C_TYPE(\"void*\")"},
+    {"FREE",       1, "C_TYPE(\"void*\")"},
 };
 
 inline const BuiltinInfo *findBuiltin(const std::string &Name) {
diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index dff4009..ffabf8d 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -1,6 +1,7 @@
 #include "cps/CodeGen.h"
 #include "cps/ArrayHandler.h"
 #include "cps/FileHandler.h"
+#include "cps/MemoryHandler.h"
 #include "cps/ExceptionHandler.h"
 #include "Builtins.h"
 #include "llvm/IR/Instructions.h"
@@ -61,6 +62,9 @@ CodeGen::CodeGen() {
     Files = std::make_unique<FileHandler>(*TheContext, *Builder, *TheModule, *RuntimeChecker,
                                          *Exceptions);
 
+    Memory = std::make_unique<MemoryHandler>(*TheContext, *Builder, *TheModule, *Types,
+                                             *RuntimeChecker, *Exceptions);
+
     // Array parameters are bound into ArrayHandler through this inversion so
     // FunctionGen never sees the array subsystem.
     FuncGen->setArrayParamBinder(
diff --git a/lib/CodeGen/ExpressionGen.cc b/lib/CodeGen/ExpressionGen.cc
index fa2581f..0a24c8d 100644
--- a/lib/CodeGen/ExpressionGen.cc
+++ b/lib/CodeGen/ExpressionGen.cc
@@ -1,9 +1,12 @@
 #include "cps/CodeGen.h"
 #include "cps/ArrayHandler.h"
 #include "cps/FileHandler.h"
+#include "cps/MemoryHandler.h"
 #include "cps/Lexer.h"
 #include "cps/TypeAST.h"
 #include "Builtins.h"
+#include <climits>
+#include <cstddef>
 
 using namespace llvm;
 using namespace cps;
@@ -184,6 +187,16 @@ Value *CodeGen::emitCall(const std::string &Callee,
             return nullptr;
         }
     }
+    if (Name == "MALLOC") {
+        Value *Size = emitAllocationSize(ArgExprs[0].get());
+        if (!Size) return nullptr;
+        return Memory->emitAllocate(Size, Line);
+    }
+    if (Name == "FREE") {
+        Value *Address = emitCoercedExpr(ArgExprs[0].get(), resolveType("C_TYPE(\"void*\")"));
+        if (!Address) return nullptr;
+        return Memory->emitRelease(Address, Line);
+    }
     if (Name == "LENGTH") {
         Value *Str = emitCoercedExpr(ArgExprs[0].get(), resolveType("STRING"));
         if (!Str) return nullptr;
@@ -279,3 +292,19 @@ Value *CodeGen::emitCall(const std::string &Callee,
     if (!marshalCallArgs(Name, ArgExprs, Line, Args)) return nullptr;
     return FuncGen->emitCall(Name, Args);
 }
+
+Value *CodeGen::emitAllocationSize(ExprAST *Expr) {
+    SaveAndRestore<int> Location(DiagnosticLine, Expr->getLine() ? Expr->getLine() : DiagnosticLine);
+    const TypeInfo *Info = getExprTypeInfo(Expr);
+    if (isWholeArrayVar(Expr) || !Info ||
+        !(Info->Kind == TypeKind::Integer ||
+          (Info->Kind == TypeKind::CScalar && Info->isIntegral()))) {
+        reportError("Allocation size must be INTEGER or a C integer");
+        return nullptr;
+    }
+    Value *Size = emitExpr(Expr);
+    if (!Size) return nullptr;
+    auto *SizeTy = Type::getIntNTy(*TheContext, sizeof(std::size_t) * CHAR_BIT);
+    return Builder->CreateIntCast(Size, SizeTy,
+                                  Info->Kind == TypeKind::Integer || Info->CSigned, "allocation_size");
+}
diff --git a/lib/CodeGen/MemoryHandler.cc b/lib/CodeGen/MemoryHandler.cc
new file mode 100644
index 0000000..fc6d982
--- /dev/null
+++ b/lib/CodeGen/MemoryHandler.cc
@@ -0,0 +1,120 @@
+#include "cps/MemoryHandler.h"
+#include "cps/ExceptionHandler.h"
+#include "cps/RuntimeCheck.h"
+#include "cps/TypeSystem.h"
+#include <climits>
+#include <cstddef>
+
+using namespace llvm;
+using namespace cps;
+
+Function *MemoryHandler::createHelper(const char *Name, ArrayRef<Type *> Args) {
+    auto *F = Function::Create(FunctionType::get(PointerType::getUnqual(Context), Args, false),
+                               Function::InternalLinkage, Name, Module);
+    Exceptions.prepareFunction(F);
+    Builder.SetInsertPoint(BasicBlock::Create(Context, "entry", F));
+    return F;
+}
+
+void MemoryHandler::ensureRuntime() {
+    if (AllocateFn) return;
+    IRBuilderBase::InsertPointGuard Guard(Builder);
+    auto *Ptr = PointerType::getUnqual(Context);
+    auto *I32 = Type::getInt32Ty(Context);
+    SizeTy = Type::getIntNTy(Context, sizeof(std::size_t) * CHAR_BIT);
+    NodeTy = StructType::create(Context, {Ptr, Ptr}, "__cps_memory_node");
+    Head = new GlobalVariable(Module, Ptr, false, GlobalValue::InternalLinkage,
+                              ConstantPointerNull::get(Ptr), "__cps_memory_head");
+    Malloc = Module.getOrInsertFunction("malloc", FunctionType::get(Ptr, {SizeTy}, false));
+    Free = Module.getOrInsertFunction("free", FunctionType::get(Type::getVoidTy(Context), {Ptr}, false));
+    AllocationError = Builder.CreateGlobalString("[Fatal] line %d: Memory allocation failed\n",
+                                                  "err_allocation", 0, &Module);
+    TrackingError = Builder.CreateGlobalString("[Fatal] line %d: Allocation tracking failed\n",
+                                                "err_tracking", 0, &Module);
+    InvalidAddress = Builder.CreateGlobalString("[Fatal] line %d: Invalid memory address (not an active allocation)\n",
+                                                 "err_memory_address", 0, &Module);
+    FindFn = createHelper("__cps_memory_find", {Ptr, I32});
+    buildFind();
+    AllocateFn = createHelper("__cps_memory_allocate", {SizeTy, I32});
+    buildAllocate();
+    ReleaseFn = createHelper("__cps_memory_release", {Ptr, I32});
+    buildRelease();
+}
+
+// Return the address of the list link that owns the matching node. This also
+// lets release unlink the head or an interior node through the same operation.
+void MemoryHandler::buildFind() {
+    auto *Ptr = PointerType::getUnqual(Context);
+    Value *Address = FindFn->getArg(0), *Line = FindFn->getArg(1);
+    auto *Entry = Builder.GetInsertBlock();
+    auto *Loop = BasicBlock::Create(Context, "search", FindFn);
+    auto *Found = BasicBlock::Create(Context, "found", FindFn);
+    Builder.CreateBr(Loop);
+    Builder.SetInsertPoint(Loop);
+    auto *Link = Builder.CreatePHI(Ptr, 2, "link");
+    Link->addIncoming(Head, Entry);
+    Value *Node = Builder.CreateLoad(Ptr, Link, "node");
+    Checker.emitCheck(Builder.CreateIsNull(Node), InvalidAddress, {Line});
+    Value *Stored = Builder.CreateLoad(Ptr, Builder.CreateStructGEP(NodeTy, Node, 0));
+    Value *NextLink = Builder.CreateStructGEP(NodeTy, Node, 1);
+    Link->addIncoming(NextLink, Builder.GetInsertBlock());
+    Builder.CreateCondBr(Builder.CreateICmpEQ(Stored, Address), Found, Loop);
+    Builder.SetInsertPoint(Found);
+    Builder.CreateRet(Link);
+}
+
+void MemoryHandler::buildAllocate() {
+    auto *Ptr = PointerType::getUnqual(Context);
+    Value *Size = AllocateFn->getArg(0), *Line = AllocateFn->getArg(1);
+    auto *Zero = BasicBlock::Create(Context, "zero_size", AllocateFn);
+    auto *Allocate = BasicBlock::Create(Context, "allocate", AllocateFn);
+    Builder.CreateCondBr(Builder.CreateICmpEQ(Size, ConstantInt::get(SizeTy, 0)), Zero, Allocate);
+    Builder.SetInsertPoint(Zero);
+    Builder.CreateRet(ConstantPointerNull::get(Ptr));
+    Builder.SetInsertPoint(Allocate);
+    Value *NodeSize = Builder.CreateIntCast(Types.getSizeOfConstant(NodeTy), SizeTy, false);
+    Value *Node = Builder.CreateCall(Malloc, {NodeSize}, "new_node");
+    Checker.emitCheck(Builder.CreateIsNull(Node), TrackingError, {Line});
+    Value *Address = Builder.CreateCall(Malloc, {Size}, "allocation");
+    Value *Failed = Builder.CreateIsNull(Address);
+    auto *Cleanup = BasicBlock::Create(Context, "cleanup_node", AllocateFn);
+    auto *Checked = BasicBlock::Create(Context, "check_allocation", AllocateFn);
+    Builder.CreateCondBr(Failed, Cleanup, Checked);
+    Builder.SetInsertPoint(Cleanup);
+    Builder.CreateCall(Free, {Node});
+    Builder.CreateBr(Checked);
+    Builder.SetInsertPoint(Checked);
+    Checker.emitCheck(Failed, AllocationError, {Line});
+    Builder.CreateStore(Address, Builder.CreateStructGEP(NodeTy, Node, 0));
+    Builder.CreateStore(Builder.CreateLoad(Ptr, Head), Builder.CreateStructGEP(NodeTy, Node, 1));
+    Builder.CreateStore(Node, Head);
+    Builder.CreateRet(Address);
+}
+
+void MemoryHandler::buildRelease() {
+    auto *Ptr = PointerType::getUnqual(Context);
+    Value *Address = ReleaseFn->getArg(0), *Line = ReleaseFn->getArg(1);
+    auto *Done = BasicBlock::Create(Context, "done", ReleaseFn);
+    auto *Release = BasicBlock::Create(Context, "release", ReleaseFn);
+    Builder.CreateCondBr(Builder.CreateIsNull(Address), Done, Release);
+    Builder.SetInsertPoint(Release);
+    Value *Link = Exceptions.emitCall(FindFn, {Address, Line}, "link");
+    Value *Node = Builder.CreateLoad(Ptr, Link);
+    Value *Next = Builder.CreateLoad(Ptr, Builder.CreateStructGEP(NodeTy, Node, 1));
+    Builder.CreateStore(Next, Link);
+    Builder.CreateCall(Free, {Address});
+    Builder.CreateCall(Free, {Node});
+    Builder.CreateBr(Done);
+    Builder.SetInsertPoint(Done);
+    Builder.CreateRet(Address);
+}
+
+Value *MemoryHandler::emitAllocate(Value *Size, int Line) {
+    ensureRuntime();
+    return Exceptions.emitCall(AllocateFn, {Size, Builder.getInt32(Line)}, "allocated");
+}
+
+Value *MemoryHandler::emitRelease(Value *Address, int Line) {
+    ensureRuntime();
+    return Exceptions.emitCall(ReleaseFn, {Address, Builder.getInt32(Line)}, "released");
+}