LLVM compiler: Allocate all variables that are free in some lambda form on the heap.

2 files changed, 30 insertions, 0 deletions

diff --git a/MLKLLVMCompiler.h b/MLKLLVMCompiler.h
index d398f0f..cdcda0d 100644
--- a/MLKLLVMCompiler.h
+++ b/MLKLLVMCompiler.h
@@ -45,6 +45,8 @@ using namespace llvm;
 +(void) processTopLevelForm:(id)object
                      inMode:(enum MLKProcessingMode)mode;
 
++(void) markVariablesForHeapAllocationInForm:(MLKForm *)form;
+
 #ifdef __cplusplus
 +(Value *) processForm:(MLKForm *)form;
 
diff --git a/MLKLLVMCompiler.mm b/MLKLLVMCompiler.mm
index 071146a..17186e7 100644
--- a/MLKLLVMCompiler.mm
+++ b/MLKLLVMCompiler.mm
@@ -146,6 +146,7 @@ static Constant
   MLKForm *form = [MLKForm formWithObject:object
                            inContext:context
                            forCompiler:self];
+  [self markVariablesForHeapAllocationInForm:form];
 
   block = BasicBlock::Create ("entry", function);
   builder.SetInsertPoint (block);
@@ -206,6 +207,33 @@ static Constant
   return [form processForLLVM];
 }
 
++(void) markVariablesForHeapAllocationInForm:(MLKForm *)form
+{
+  NSArray *subforms = [form subforms];
+  unsigned int i;
+
+  for (i = 0; i < [subforms count]; i++)
+    {
+      MLKForm *subform = [subforms objectAtIndex:i];
+
+      [self markVariablesForHeapAllocationInForm:subform];
+
+      if ([subform isKindOfClass:[MLKSimpleLambdaForm class]]
+          || [subform isKindOfClass:[MLKLambdaForm class]])
+        {
+          NSArray *freeVariables = [[subform freeVariables] allObjects];
+          unsigned int j;
+
+          for (j = 0; j < [freeVariables count]; j++)
+            {
+              id variable = [freeVariables objectAtIndex:j];
+              [variable setVariableHeapAllocation:YES
+                        forSymbol:variable];
+            }
+        }
+    }
+}
+
 +(Value *) insertSelectorLookup:(NSString *)name
 {
   Constant *function =