resolve parameters, infer function return type from body expression

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

@@ -2,6 +2,9 @@ package org.jetbrains.jet.codegen;
 
 import gnu.trove.TObjectIntHashMap;
 import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.types.DeclarationDescriptor;
+import org.jetbrains.jet.lang.types.ValueParameterDescriptor;
 import org.objectweb.asm.Label;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
@@ -21,9 +24,11 @@ public class ExpressionCodegen extends JetVisitor {
     private final InstructionAdapter v;
     private final TObjectIntHashMap<JetProperty> myVarIndex = new TObjectIntHashMap<JetProperty>();
     private       int myMaxVarIndex = 0;
+    private final BindingContext bindingContext;
 
-    public ExpressionCodegen(MethodVisitor v) {
+    public ExpressionCodegen(MethodVisitor v, BindingContext bindingContext) {
         this.v = new InstructionAdapter(v);
+        this.bindingContext = bindingContext;
     }
 
     private void gen(JetElement expr) {
@@ -243,6 +248,18 @@ public class ExpressionCodegen extends JetVisitor {
         }
     }
 
+    @Override
+    public void visitReferenceExpression(JetReferenceExpression expression) {
+        final DeclarationDescriptor descriptor = bindingContext.resolve(expression);
+        if (descriptor instanceof ValueParameterDescriptor) {
+            final int index = ((ValueParameterDescriptor) descriptor).getIndex();
+            v.visitVarInsn(Opcodes.ALOAD, index);  // TODO +1 for non-static methods
+        }
+        else {
+            throw new UnsupportedOperationException("don't know how to generate reference " + descriptor);
+        }
+    }
+
     private Type getType(JetProperty var) {
         return InstructionAdapter.OBJECT_TYPE;  // TODO:
     }

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

 package org.jetbrains.jet.codegen;
 
 import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.types.FunctionDescriptor;
+import org.jetbrains.jet.lang.types.JetStandardClasses;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
@@ -14,9 +17,11 @@ import java.util.List;
  */
 public class FunctionCodegen {
     private final ClassVisitor v;
+    private final BindingContext bindingContext;
 
-    public FunctionCodegen(ClassVisitor v) {
+    public FunctionCodegen(ClassVisitor v, BindingContext bindingContext) {
         this.v = v;
+        this.bindingContext = bindingContext;
     }
 
     public void gen(JetFunction f, JetNamespace owner) {
@@ -26,28 +31,50 @@ public class FunctionCodegen {
             parameterTypes[i] = mapTypeReference(parameters.get(i).getTypeReference());
         }
         final JetTypeReference returnTypeRef = f.getReturnTypeRef();
-        Type returnType = returnTypeRef == null ? Type.VOID_TYPE : mapTypeReference(returnTypeRef);
+        Type returnType;
+        if (returnTypeRef == null) {
+            final FunctionDescriptor functionDescriptor = bindingContext.getFunctionDescriptor(f);
+            final org.jetbrains.jet.lang.types.Type type = functionDescriptor.getUnsubstitutedReturnType();
+            if (type.equals(JetStandardClasses.getUnitType())) {
+                returnType = Type.VOID_TYPE;
+            }
+            else if (type.equals(JetStandardClasses.getIntType())) {
+                returnType = Type.getType(Integer.class);
+            }
+            else {
+                throw new UnsupportedOperationException("don't know how to map type " + type);
+            }
+        }
+        else {
+            returnType = mapTypeReference(returnTypeRef);
+        }
         Method method = new Method(f.getName(), returnType, parameterTypes);
         final MethodVisitor mv = v.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
                 method.getName(), method.getDescriptor(), null, null);
         mv.visitCode();
         final JetExpression bodyExpression = f.getBodyExpression();
-        bodyExpression.accept(new ExpressionCodegen(mv));
-        if (needReturn(bodyExpression)) {
-            mv.visitInsn(Opcodes.RETURN);
-        }
+        bodyExpression.accept(new ExpressionCodegen(mv, bindingContext));
+        generateReturn(mv, bodyExpression);
         mv.visitMaxs(0, 0);
         mv.visitEnd();
     }
 
-    private boolean needReturn(JetExpression bodyExpression) {
+    private void generateReturn(MethodVisitor mv, JetExpression bodyExpression) {
+        if (!endsWithReturn(bodyExpression)) {
+            final org.jetbrains.jet.lang.types.Type expressionType = bindingContext.getExpressionType(bodyExpression);
+            if (expressionType.equals(JetStandardClasses.getUnitType())) {
+                mv.visitInsn(Opcodes.RETURN);
+            }
+            else {
+                mv.visitInsn(Opcodes.ARETURN);
+            }
+        }
+    }
+
+    private static boolean endsWithReturn(JetExpression bodyExpression) {
         if (bodyExpression instanceof JetBlockExpression) {
             final List<JetElement> statements = ((JetBlockExpression) bodyExpression).getStatements();
-            if (statements.size() == 0) {
-                return true;
-            }
-            final JetElement jetElement = statements.get(statements.size() - 1);
-            return !(jetElement instanceof JetReturnExpression);
+            return statements.size() > 0 && statements.get(statements.size()-1) instanceof JetReturnExpression;
         }
         return false;
     }
@@ -58,26 +85,30 @@ public class FunctionCodegen {
         }
         final JetTypeElement typeElement = typeRef.getTypeElement();
         if (typeElement instanceof JetUserType) {
-            final String referencedName = ((JetUserType) typeElement).getReferencedName();
-            if ("Array".equals(referencedName)) {
-                final List<JetTypeProjection> typeArguments = ((JetUserType) typeElement).getTypeArguments();
-                if (typeArguments.size() != 1) {
-                    throw new UnsupportedOperationException("arrays must have one type argument");
-                }
-                final JetTypeReference elementTypeRef = typeArguments.get(0).getTypeReference();
-                Type elementType = mapTypeReference(elementTypeRef);
-                return Type.getType("[" + elementType.getDescriptor());
-            }
+            final JetUserType userType = (JetUserType) typeElement;
+            return mapType(userType.getReferencedName(), userType.getTypeArguments());
+        }
 
-            if ("String".equals(referencedName)) {
-                return Type.getType(String.class);
-            }
-            if ("Int".equals(referencedName)) {
-                return Type.getType(Integer.class);
+        throw new UnsupportedOperationException("Unknown type " + typeRef);
+    }
+
+    private Type mapType(final String name, final List<JetTypeProjection> typeArguments) {
+        if ("Array".equals(name)) {
+            if (typeArguments.size() != 1) {
+                throw new UnsupportedOperationException("arrays must have one type argument");
             }
+            final JetTypeReference elementTypeRef = typeArguments.get(0).getTypeReference();
+            Type elementType = mapTypeReference(elementTypeRef);
+            return Type.getType("[" + elementType.getDescriptor());
         }
 
-        throw new UnsupportedOperationException("Unknown type " + typeRef);
+        if ("String".equals(name)) {
+            return Type.getType(String.class);
+        }
+        if ("Int".equals(name)) {
+            return Type.getType(Integer.class);
+        }
+        throw new UnsupportedOperationException("Unknown type " + name);
     }
 
     public void gen(JetFunction f, JetClass owner) {

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

@@ -4,9 +4,14 @@ import org.jetbrains.jet.lang.psi.JetDeclaration;
 import org.jetbrains.jet.lang.psi.JetFunction;
 import org.jetbrains.jet.lang.psi.JetNamespace;
 import org.jetbrains.jet.lang.psi.JetProperty;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.resolve.TopDownAnalyzer;
+import org.jetbrains.jet.lang.types.JetStandardClasses;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.Opcodes;
 
+import java.util.List;
+
 /**
  * @author max
  */
@@ -15,8 +20,11 @@ public class NamespaceCodegen {
     }
 
     public void generate(JetNamespace namespace, ClassVisitor v) {
+        List<JetDeclaration> declarations = namespace.getDeclarations();
+        BindingContext bindingContext = new TopDownAnalyzer().process(JetStandardClasses.STANDARD_CLASSES, declarations);
+
         final PropertyCodegen propertyCodegen = new PropertyCodegen(v);
-        final FunctionCodegen functionCodegen = new FunctionCodegen(v);
+        final FunctionCodegen functionCodegen = new FunctionCodegen(v, bindingContext);
         v.visit(Opcodes.V1_6,
                 Opcodes.ACC_PUBLIC,
                 getJVMClassName(namespace),

idea/src/org/jetbrains/jet/lang/types/ValueParameterDescriptor.java

@@ -4,6 +4,12 @@ package org.jetbrains.jet.lang.types;
  * @author abreslav
  */
 public interface ValueParameterDescriptor extends PropertyDescriptor {
+    /**
+     * Returns the 0-based index of the value parameter in the parameter list of its containing function.
+     *
+     * @return the parameter index
+     */
+    int getIndex();
     boolean hasDefaultValue();
     boolean isRef();
     boolean isVararg();

idea/src/org/jetbrains/jet/lang/types/ValueParameterDescriptorImpl.java

 package org.jetbrains.jet.lang.types;
 
+import com.intellij.psi.PsiElement;
+import org.jetbrains.jet.lang.psi.JetDeclaration;
 import org.jetbrains.jet.lang.psi.JetParameter;
+import org.jetbrains.jet.lang.psi.JetParameterList;
 
 import java.util.List;
 
@@ -17,6 +20,16 @@ public class ValueParameterDescriptorImpl extends PropertyDescriptorImpl impleme
         this.isVararg = isVararg;
     }
 
+    @Override
+    public int getIndex() {
+        final JetDeclaration element = getPsiElement();
+        final PsiElement parent = element.getParent();
+        if (parent instanceof JetParameterList) {
+            return ((JetParameterList) parent).getParameters().indexOf(element);
+        }
+        throw new IllegalStateException("couldn't find index for parameter");
+    }
+
     @Override
     public boolean hasDefaultValue() {
         return hasDefaultValue;

idea/testData/codegen/returnA.jet

+fun foo(a : Int) = a

idea/tests/org/jetbrains/jet/codegen/ExpressionGenTest.java

@@ -60,7 +60,7 @@ public class ExpressionGenTest extends LightDaemonAnalyzerTestCase {
         JetProperty p = PsiTreeUtil.getParentOfType(getFile().findElementAt(0), JetProperty.class);
 
         TraceMethodVisitor trace = new TraceMethodVisitor();
-        p.getInitializer().accept(new ExpressionCodegen(trace));
+        p.getInitializer().accept(new ExpressionCodegen(trace, null));
 
         StringWriter out = new StringWriter();
         trace.print(new PrintWriter(out));

idea/tests/org/jetbrains/jet/codegen/NamespaceGenTest.java

@@ -45,6 +45,17 @@ public class NamespaceGenTest extends LightCodeInsightFixtureTestCase {
         assertEquals(new Integer(42), returnValue);
     }
 
+    public void testReturnA() throws Exception {
+        myFixture.configureByFile(JetParsingTest.getTestDataDir() + "/codegen/returnA.jet");
+        final String text = generateToText();
+        System.out.println(text);
+
+        final Class aClass = generateToClass();
+        final Method main = firstMethod(aClass);
+        final Object returnValue = main.invoke(null, 42);
+        assertEquals(new Integer(42), returnValue);
+    }
+
     private String generateToText() {
         StringWriter writer = new StringWriter();
         JetFile jetFile = (JetFile) myFixture.getFile();