first 'real' closure test, green

idea/src/org/jetbrains/jet/codegen/ClosureCodegen.java

@@ -3,11 +3,12 @@
  */
 package org.jetbrains.jet.codegen;
 
-import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Pair;
 import com.intellij.psi.util.PsiTreeUtil;
+import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
 import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
 import org.jetbrains.jet.lang.descriptors.ValueParameterDescriptor;
+import org.jetbrains.jet.lang.descriptors.VariableDescriptor;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.types.JetType;
 import org.objectweb.asm.ClassVisitor;
@@ -19,13 +20,21 @@ import org.objectweb.asm.commons.Method;
 import org.objectweb.asm.signature.SignatureWriter;
 
 import java.util.Arrays;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 
 public class ClosureCodegen {
     private final GenerationState state;
+    private final ExpressionCodegen context;
+    private ClassVisitor cv = null;
+    private String name = null;
 
-    public ClosureCodegen(GenerationState state) {
+    private Map<DeclarationDescriptor, EnclosedValueDescriptor> closure = new LinkedHashMap<DeclarationDescriptor, EnclosedValueDescriptor>();
+
+    public ClosureCodegen(GenerationState state, ExpressionCodegen context) {
         this.state = state;
+        this.context = context;
     }
 
     public static Method erasedInvokeSignature(FunctionDescriptor fd) {
@@ -44,6 +53,31 @@ public class ClosureCodegen {
         return state.getTypeMapper().mapSignature("invoke", fd);
     }
 
+    public StackValue lookupInContext(DeclarationDescriptor d) {
+        if (d instanceof VariableDescriptor) {
+            VariableDescriptor vd = (VariableDescriptor) d;
+
+            EnclosedValueDescriptor answer = closure.get(vd);
+            if (answer != null) return answer.getInnerValue();
+
+            final int idx = context.lookupLocal(vd);
+            if (idx < 0) return null;
+
+            final Type type = state.getTypeMapper().mapType(vd.getOutType());
+
+            StackValue outerValue = StackValue.local(idx, type);
+            final String fieldName = "$" + (closure.size() + 1);
+            StackValue innerValue = StackValue.field(type, name, fieldName, false);
+            cv.visitField(Opcodes.ACC_PUBLIC, fieldName, type.getDescriptor(), null, null);
+            answer = new EnclosedValueDescriptor(d, innerValue, outerValue);
+            closure.put(d, answer);
+
+            return innerValue;
+        }
+
+        return null;
+    }
+
     public GeneratedClosureDescriptor gen(JetFunctionLiteralExpression fun) {
         JetNamedDeclaration container = PsiTreeUtil.getParentOfType(fun, JetNamespace.class, JetClass.class, JetObjectDeclaration.class);
 
@@ -58,8 +92,8 @@ public class ClosureCodegen {
 
         final FunctionDescriptor funDescriptor = (FunctionDescriptor) state.getBindingContext().getDeclarationDescriptor(fun);
 
-        final ClassVisitor cv = nameAndVisitor.getSecond();
+        cv = nameAndVisitor.getSecond();
-        final String name = nameAndVisitor.getFirst();
+        name = nameAndVisitor.getFirst();
 
         SignatureWriter signatureWriter = new SignatureWriter();
 
@@ -79,19 +113,25 @@ public class ClosureCodegen {
                 null,
                 funClass,
                 new String[0]
-                );
+        );
 
-        final Method constructor = generateConstructor(cv, funClass);
 
         generateBridge(name, funDescriptor, cv);
-        generateBody(name, funDescriptor, cv, container.getProject(), fun.getFunctionLiteral().getBodyExpression().getStatements());
+        generateBody(funDescriptor, cv, fun.getFunctionLiteral().getBodyExpression().getStatements());
+
+        final Method constructor = generateConstructor(funClass);
 
         cv.visitEnd();
 
-        return new GeneratedClosureDescriptor(name, constructor);
+        final GeneratedClosureDescriptor answer = new GeneratedClosureDescriptor(name, constructor);
+        for (DeclarationDescriptor descriptor : closure.keySet()) {
+            final EnclosedValueDescriptor valueDescriptor = closure.get(descriptor);
+            answer.addArg(valueDescriptor.getOuterValue());
+        }
+        return answer;
     }
 
-    private void generateBody(String className, FunctionDescriptor funDescriptor, ClassVisitor cv, Project project, List<JetElement> body) {
+    private void generateBody(FunctionDescriptor funDescriptor, ClassVisitor cv, List<JetElement> body) {
         FunctionCodegen fc = new FunctionCodegen(null, cv, state);
         fc.generatedMethod(body, OwnerKind.IMPLEMENTATION, invokeSignature(funDescriptor), funDescriptor.getReceiverType(), funDescriptor.getUnsubstitutedValueParameters(), null);
     }
@@ -131,8 +171,14 @@ public class ClosureCodegen {
         mv.visitEnd();
     }
 
-    private Method generateConstructor(ClassVisitor cv, String funClass) {
+    private Method generateConstructor(String funClass) {
-        final Method constructor = new Method("<init>", Type.VOID_TYPE, new Type[0]); // TODO:
+        Type[] argTypes = new Type[closure.size()];
+        int i = 0;
+        for (DeclarationDescriptor descriptor : closure.keySet()) {
+            argTypes[i++] = state.getTypeMapper().mapType(((VariableDescriptor) descriptor).getOutType());
+        }
+
+        final Method constructor = new Method("<init>", Type.VOID_TYPE, argTypes);
         final MethodVisitor mv = cv.visitMethod(Opcodes.ACC_PUBLIC, "<init>", constructor.getDescriptor(), null, new String[0]);
         mv.visitCode();
         InstructionAdapter iv = new InstructionAdapter(mv);
@@ -140,8 +186,14 @@ public class ClosureCodegen {
         iv.load(0, Type.getObjectType(funClass));
         iv.invokespecial(funClass, "<init>", "()V");
 
-        iv.visitInsn(Opcodes.RETURN);
+        for (int j = 0; j < argTypes.length; j++) {
+            Type type = argTypes[j];
+            StackValue.local(0, JetTypeMapper.TYPE_OBJECT).put(JetTypeMapper.TYPE_OBJECT, iv);
+            StackValue.local(j + 1, type).put(type, iv);
+            StackValue.field(type, name, "$" + (j + 1), false).store(iv);
+        }
 
+        iv.visitInsn(Opcodes.RETURN);
 
         mv.visitMaxs(0, 0);
         mv.visitEnd();

idea/src/org/jetbrains/jet/codegen/EnclosedValueDescriptor.java

@@ -0,0 +1,30 @@
+/*
+ * @author max
+ */
+package org.jetbrains.jet.codegen;
+
+import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
+
+public class EnclosedValueDescriptor {
+    private final DeclarationDescriptor descriptor;
+    private final StackValue innerValue;
+    private final StackValue outerValue;
+
+    public EnclosedValueDescriptor(DeclarationDescriptor descriptor, StackValue innerValue, StackValue outerValue) {
+        this.descriptor = descriptor;
+        this.innerValue = innerValue;
+        this.outerValue = outerValue;
+    }
+
+    public DeclarationDescriptor getDescriptor() {
+        return descriptor;
+    }
+
+    public StackValue getInnerValue() {
+        return innerValue;
+    }
+
+    public StackValue getOuterValue() {
+        return outerValue;
+    }
+}

idea/src/org/jetbrains/jet/codegen/ExpressionCodegen.java

@@ -289,7 +289,7 @@ public class ExpressionCodegen extends JetVisitor {
         if (asmParamType.getSort() == Type.OBJECT && !"java.lang.Object".equals(asmParamType.getClassName())) {
             v.checkcast(asmParamType);
         }
-        v.store(myMap.getIndex(parameterDescriptor), asmParamType);
+        v.store(lookupLocal(parameterDescriptor), asmParamType);
 
         gen(expression.getBody(), Type.VOID_TYPE);
 
@@ -384,7 +384,7 @@ public class ExpressionCodegen extends JetVisitor {
             v.load(myIndexVar, Type.INT_TYPE);
             v.aload(loopRangeType.getElementType());
             StackValue.onStack(loopRangeType.getElementType()).put(asmParamType, v);
-            v.store(myMap.getIndex(parameterDescriptor), asmParamType);
+            v.store(lookupLocal(parameterDescriptor), asmParamType);
         }
 
         protected void generateIncrement() {
@@ -414,21 +414,21 @@ public class ExpressionCodegen extends JetVisitor {
             v.store(myRangeVar, loopRangeType);
 
             v.invokevirtual("jet/IntRange", "getStartValue", "()I");
-            v.store(myMap.getIndex(parameterDescriptor), Type.INT_TYPE);
+            v.store(lookupLocal(parameterDescriptor), Type.INT_TYPE);
             v.invokevirtual("jet/IntRange", "getEndValue", "()I");
             v.store(myEndVar, Type.INT_TYPE);
         }
 
         @Override
         protected void generateCondition(Type asmParamType, Label end) {
-            v.load(myMap.getIndex(parameterDescriptor), Type.INT_TYPE);
+            v.load(lookupLocal(parameterDescriptor), Type.INT_TYPE);
             v.load(myEndVar, Type.INT_TYPE);
             v.ificmpgt(end);
         }
 
         @Override
         protected void generateIncrement() {
-            v.iinc(myMap.getIndex(parameterDescriptor), 1);  // TODO support decreasing order
+            v.iinc(lookupLocal(parameterDescriptor), 1);  // TODO support decreasing order
         }
 
         @Override
@@ -482,10 +482,19 @@ public class ExpressionCodegen extends JetVisitor {
             generateBlock(expression.getFunctionLiteral().getBodyExpression().getStatements());
         }
         else {
-            final GeneratedClosureDescriptor closure = new ClosureCodegen(state).gen(expression);
+            final GeneratedClosureDescriptor closure = state.generateClosure(expression, this);
+
             v.anew(Type.getObjectType(closure.getClassname()));
             v.dup();
-            v.invokespecial(closure.getClassname(), "<init>", closure.getConstructor().getDescriptor());
+
+            final Method cons = closure.getConstructor();
+
+            for (int i = 0; i < closure.getArgs().size(); i++) {
+                StackValue arg = closure.getArgs().get(i);
+                arg.put(cons.getArgumentTypes()[i], v);
+            }
+
+            v.invokespecial(closure.getClassname(), "<init>", cons.getDescriptor());
         }
     }
 
@@ -548,6 +557,7 @@ public class ExpressionCodegen extends JetVisitor {
     @Override
     public void visitSimpleNameExpression(JetSimpleNameExpression expression) {
         DeclarationDescriptor descriptor = bindingContext.resolveReferenceExpression(expression);
+        if (descriptor instanceof NamespaceDescriptor) return; // No code to generate
 
         if (descriptor instanceof VariableAsFunctionDescriptor) {
             descriptor = ((VariableAsFunctionDescriptor) descriptor).getVariableDescriptor();
@@ -577,7 +587,7 @@ public class ExpressionCodegen extends JetVisitor {
             myStack.push(StackValue.field(fieldType, owner, psiField.getName(), isStatic));
         }
         else {
-            int index = myMap.getIndex(descriptor);
+            int index = lookupLocal(descriptor);
             if (index >= 0) {
                 final JetType outType = ((VariableDescriptor) descriptor).getOutType();
                 myStack.push(StackValue.local(index, typeMapper.mapType(outType)));
@@ -594,7 +604,7 @@ public class ExpressionCodegen extends JetVisitor {
                         for (ValueParameterDescriptor parameter : parameters) {
                             if (parameter.getName().equals(descriptor.getName())) {
                                 final JetType outType = ((VariableDescriptor) descriptor).getOutType();
-                                myStack.push(StackValue.local(myMap.getIndex(parameter), typeMapper.mapType(outType)));
+                                myStack.push(StackValue.local(lookupLocal(parameter), typeMapper.mapType(outType)));
                                 return;
                             }
                         }
@@ -626,12 +636,22 @@ public class ExpressionCodegen extends JetVisitor {
                     myStack.push(iValue);
                 }
             }
-            else if (!(descriptor instanceof NamespaceDescriptor)) {
+            else {
-                throw new UnsupportedOperationException("don't know how to generate reference " + descriptor);
+                final StackValue value = state.lookupInContext(descriptor);
+                if (value == null) {
+                    throw new UnsupportedOperationException("don't know how to generate reference " + descriptor);
+                }
+                // receiver
+                StackValue.local(0, JetTypeMapper.TYPE_OBJECT).put(JetTypeMapper.TYPE_OBJECT, v);
+                myStack.push(value);
             }
         }
     }
 
+    public int lookupLocal(DeclarationDescriptor descriptor) {
+        return myMap.getIndex(descriptor);
+    }
+
     public StackValue intermediateValueForProperty(PropertyDescriptor propertyDescriptor, final boolean forceField, boolean forceInterface) {
         DeclarationDescriptor containingDeclaration = propertyDescriptor.getContainingDeclaration();
         boolean isStatic = containingDeclaration instanceof NamespaceDescriptorImpl;
@@ -890,7 +910,7 @@ public class ExpressionCodegen extends JetVisitor {
 
     private int indexOfLocal(JetReferenceExpression lhs) {
         final DeclarationDescriptor declarationDescriptor = bindingContext.resolveReferenceExpression(lhs);
-        return myMap.getIndex(declarationDescriptor);
+        return lookupLocal(declarationDescriptor);
     }
 
     private static Type psiTypeToAsm(PsiType type) {
@@ -1369,7 +1389,7 @@ public class ExpressionCodegen extends JetVisitor {
     @Override
     public void visitProperty(JetProperty property) {
         VariableDescriptor variableDescriptor = bindingContext.getVariableDescriptor(property);
-        int index = myMap.getIndex(variableDescriptor);
+        int index = lookupLocal(variableDescriptor);
 
         assert index >= 0;
 
@@ -1538,7 +1558,7 @@ public class ExpressionCodegen extends JetVisitor {
             VariableDescriptor descriptor = bindingContext.getVariableDescriptor(clause.getCatchParameter());
             Type descriptorType = typeMapper.mapType(descriptor.getOutType());
             myMap.enter(descriptor, 1);
-            int index = myMap.getIndex(descriptor);
+            int index = lookupLocal(descriptor);
             v.store(index, descriptorType);
 
             gen(clause.getCatchBody(), Type.VOID_TYPE);

idea/src/org/jetbrains/jet/codegen/GeneratedClosureDescriptor.java

@@ -5,9 +5,13 @@ package org.jetbrains.jet.codegen;
 
 import org.objectweb.asm.commons.Method;
 
+import java.util.ArrayList;
+import java.util.List;
+
 public class GeneratedClosureDescriptor {
     private final String classname;
     private Method constructor;
+    private List<StackValue> args = new ArrayList<StackValue>();
 
     public GeneratedClosureDescriptor(String classname, Method constructor) {
         this.classname = classname;
@@ -21,4 +25,12 @@ public class GeneratedClosureDescriptor {
     public Method getConstructor() {
         return constructor;
     }
+
+    public void addArg(StackValue local) {
+        args.add(local);
+    }
+
+    public List<StackValue> getArgs() {
+        return args;
+    }
 }

idea/src/org/jetbrains/jet/codegen/GenerationState.java

@@ -9,7 +9,9 @@ import com.intellij.util.containers.Stack;
 import org.jetbrains.jet.lang.ErrorHandler;
 import org.jetbrains.jet.lang.cfg.pseudocode.JetControlFlowDataTraceFactory;
 import org.jetbrains.jet.lang.descriptors.ClassDescriptor;
+import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
 import org.jetbrains.jet.lang.psi.JetFile;
+import org.jetbrains.jet.lang.psi.JetFunctionLiteralExpression;
 import org.jetbrains.jet.lang.psi.JetNamespace;
 import org.jetbrains.jet.lang.resolve.AnalyzingUtils;
 import org.jetbrains.jet.lang.resolve.BindingContext;
@@ -21,12 +23,12 @@ public class GenerationState {
     private final Project project;
 
     private JetTypeMapper typeMapper;
-    private final Stack<BindingContext> bindingContexts;
+    private final Stack<BindingContext> bindingContexts = new Stack<BindingContext>();
+    private final Stack<ClosureCodegen> closureContexts = new Stack<ClosureCodegen>();
     private final JetStandardLibrary standardLibrary;
 
     public GenerationState(Project project, boolean text) {
         this.project = project;
-        this.bindingContexts = new Stack<BindingContext>();
         this.standardLibrary = JetStandardLibrary.getJetStandardLibrary(project);
         this.factory = new ClassFileFactory(project, text, this);
     }
@@ -94,4 +96,21 @@ public class GenerationState {
             typeMapper = null;
         }
     }
+
+    public GeneratedClosureDescriptor generateClosure(JetFunctionLiteralExpression literal, ExpressionCodegen context) {
+        final ClosureCodegen codegen = new ClosureCodegen(this, context);
+        closureContexts.push(codegen);
+        try {
+            return codegen.gen(literal);
+        }
+        finally {
+            final ClosureCodegen pooped = closureContexts.pop();
+            assert pooped == codegen;
+        }
+    }
+
+    public StackValue lookupInContext(DeclarationDescriptor d) {
+        final ClosureCodegen top = closureContexts.peek();
+        return top != null ? top.lookupInContext(d) : null;
+    }
 }

idea/testData/codegen/classes/enclosingLocalVariable.jet

@@ -1,6 +1,6 @@
 fun box() : String {
     val cl = 39
-    return if (sum(200, { cl }) == 2) "OK" else "FAIL"
+    return if (sum(200, { cl }) == 239) "OK" else "FAIL"
 }
 
 fun sum(arg:Int, f : {() : Int}) : Int {