initial support for generating extension functions

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

@@ -175,7 +175,7 @@ public class ClassCodegen {
         frameMap.enterTemp();   // this
 
         final InstructionAdapter iv = new InstructionAdapter(mv);
-        ExpressionCodegen codegen = new ExpressionCodegen(mv, bindingContext, frameMap, typeMapper, Type.VOID_TYPE, classDescriptor, kind);
+        ExpressionCodegen codegen = new ExpressionCodegen(mv, bindingContext, frameMap, typeMapper, null, Type.VOID_TYPE, classDescriptor, kind);
 
         String classname = typeMapper.jvmName(classDescriptor, kind);
         final Type classType = Type.getType("L" + classname + ";");

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

@@ -66,6 +66,7 @@ public class ExpressionCodegen extends JetVisitor {
     private final InstructionAdapter v;
     private final FrameMap myMap;
     private final JetTypeMapper typeMapper;
+    private final JetType receiverType;
     private final Type returnType;
     private final DeclarationDescriptor contextType;
     private final OwnerKind contextKind;
@@ -75,11 +76,13 @@ public class ExpressionCodegen extends JetVisitor {
                              BindingContext bindingContext,
                              FrameMap myMap,
                              JetTypeMapper typeMapper,
+                             JetType receiverType,
                              Type returnType,
                              DeclarationDescriptor contextType,
                              OwnerKind contextKind) {
         this.myMap = myMap;
         this.typeMapper = typeMapper;
+        this.receiverType = receiverType;
         this.returnType = returnType;
         this.contextType = contextType;
         this.contextKind = contextKind;
@@ -107,11 +110,15 @@ public class ExpressionCodegen extends JetVisitor {
         int oldStackDepth = myStack.size();
         gen(expr);
         if (myStack.size() == oldStackDepth+1) {
-            StackValue value = myStack.pop();
+            putTopOfStack(type);
-            value.put(type, v);
         }
     }
 
+    private void putTopOfStack(Type type) {
+        StackValue value = myStack.pop();
+        value.put(type, v);
+    }
+
     public void genToJVMStack(JetExpression expr) {
         gen(expr, expressionType(expr));
     }
@@ -526,8 +533,7 @@ public class ExpressionCodegen extends JetVisitor {
 
     public void returnTopOfStack() {
         if (myStack.size() > 0) {
-            StackValue value = myStack.pop();
+            putTopOfStack(returnType);
-            value.put(returnType, v);
             v.areturn(returnType);
         }
     }
@@ -1354,23 +1360,37 @@ public class ExpressionCodegen extends JetVisitor {
 
     @Override
     public void visitThisExpression(JetThisExpression expression) {
-        thisToStack();
+        generateThis();
     }
 
     public void thisToStack() {
-        if (contextKind == OwnerKind.NAMESPACE) {
+        generateThis();
-            throw new UnsupportedOperationException("Cannot generate this expression in top level context");
+        putTopOfStack(JetTypeMapper.TYPE_OBJECT);
-        }
+    }
 
-        ClassDescriptor contextClass = (ClassDescriptor) contextType;
+    private void generateThis() {
-        if (contextKind == OwnerKind.IMPLEMENTATION) {
+        if (contextKind == OwnerKind.NAMESPACE) {
-            v.load(0, JetTypeMapper.jetImplementationType(contextClass));
+            if (receiverType != null) {
-        }
+                myStack.push(StackValue.local(0, typeMapper.mapType(receiverType)));
-        else if (contextKind == OwnerKind.DELEGATING_IMPLEMENTATION) {
+            }
-            v.getfield(typeMapper.jvmName(contextClass, contextKind), "$this", JetTypeMapper.jetInterfaceType(contextClass).getDescriptor());
+            else {
+                throw new UnsupportedOperationException("Cannot generate this expression in top level context");
+            }
         }
         else {
-            throw new UnsupportedOperationException("Unknown kind: " + contextKind);
+            ClassDescriptor contextClass = (ClassDescriptor) contextType;
+            if (contextKind == OwnerKind.IMPLEMENTATION) {
+                myStack.push(StackValue.local(0, JetTypeMapper.jetImplementationType(contextClass)));
+            }
+            else if (contextKind == OwnerKind.DELEGATING_IMPLEMENTATION) {
+                myStack.push(StackValue.field(JetTypeMapper.jetInterfaceType(contextClass),
+                        typeMapper.jvmName(contextClass, contextKind),
+                        "$this",
+                        false));
+            }
+            else {
+                throw new UnsupportedOperationException("Unknown kind: " + contextKind);
+            }
         }
     }
 

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

@@ -1,10 +1,12 @@
 package org.jetbrains.jet.codegen;
 
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
 import org.jetbrains.jet.lang.descriptors.ValueParameterDescriptor;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.types.JetStandardLibrary;
+import org.jetbrains.jet.lang.types.JetType;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
@@ -31,12 +33,18 @@ public class FunctionCodegen {
     }
 
     public void gen(JetFunction f, OwnerKind kind) {
+        final JetTypeReference receiverTypeRef = f.getReceiverTypeRef();
+        final JetType receiverType = receiverTypeRef == null ? null : bindingContext.resolveTypeReference(receiverTypeRef);
         Method method = typeMapper.mapSignature(f);
         List<ValueParameterDescriptor> paramDescrs = bindingContext.getFunctionDescriptor(f).getUnsubstitutedValueParameters();
-        generateMethod(f, kind, method, paramDescrs);
+        generateMethod(f, kind, method, receiverType, paramDescrs);
     }
 
-    public void generateMethod(JetDeclarationWithBody f, OwnerKind kind, Method jvmSignature, List<ValueParameterDescriptor> paramDescrs) {
+    public void generateMethod(JetDeclarationWithBody f,
+                               OwnerKind kind,
+                               Method jvmSignature,
+                               @Nullable JetType receiverType,
+                               List<ValueParameterDescriptor> paramDescrs) {
         int flags = Opcodes.ACC_PUBLIC; // TODO.
 
         boolean isStatic = kind == OwnerKind.NAMESPACE;
@@ -55,7 +63,7 @@ public class FunctionCodegen {
             mv.visitCode();
             FrameMap frameMap = new FrameMap();
 
-            if (kind != OwnerKind.NAMESPACE) {
+            if (kind != OwnerKind.NAMESPACE || receiverType != null) {
                 frameMap.enterTemp();  // 0 slot for this
             }
 
@@ -65,7 +73,8 @@ public class FunctionCodegen {
                 frameMap.enter(parameter, argTypes[i].getSize());
             }
 
-            ExpressionCodegen codegen = new ExpressionCodegen(mv, bindingContext, frameMap, typeMapper, jvmSignature.getReturnType(), contextDescriptor, kind);
+            ExpressionCodegen codegen = new ExpressionCodegen(mv, bindingContext, frameMap, typeMapper, receiverType,
+                    jvmSignature.getReturnType(), contextDescriptor, kind);
             if (kind instanceof OwnerKind.DelegateKind) {
                 OwnerKind.DelegateKind dk = (OwnerKind.DelegateKind) kind;
                 InstructionAdapter iv = new InstructionAdapter(mv);

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

@@ -206,10 +206,15 @@ public class JetTypeMapper {
     }
 
     public Method mapSignature(JetFunction f) {
+        final JetTypeReference receiverTypeRef = f.getReceiverTypeRef();
+        final JetType receiverType = receiverTypeRef == null ? null : bindingContext.resolveTypeReference(receiverTypeRef);
         final List<JetParameter> parameters = f.getValueParameters();
-        Type[] parameterTypes = new Type[parameters.size()];
+        List<Type> parameterTypes = new ArrayList<Type>();
-        for (int i = 0; i < parameters.size(); i++) {
+        if (receiverType != null) {
-            parameterTypes[i] = mapType(bindingContext.resolveTypeReference(parameters.get(i).getTypeReference()));
+            parameterTypes.add(mapType(receiverType));
+        }
+        for (JetParameter parameter : parameters) {
+            parameterTypes.add(mapType(bindingContext.resolveTypeReference(parameter.getTypeReference())));
         }
         final JetTypeReference returnTypeRef = f.getReturnTypeRef();
         Type returnType;
@@ -221,7 +226,7 @@ public class JetTypeMapper {
         else {
             returnType = mapType(bindingContext.resolveTypeReference(returnTypeRef));
         }
-        return new Method(f.getName(), returnType, parameterTypes);
+        return new Method(f.getName(), returnType, parameterTypes.toArray(new Type[parameterTypes.size()]));
     }
 
     public Method mapGetterSignature(PropertyDescriptor descriptor) {

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

@@ -13,7 +13,6 @@ import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.commons.InstructionAdapter;
-import org.objectweb.asm.commons.Method;
 
 /**
  * @author max
@@ -72,7 +71,7 @@ public class NamespaceCodegen {
 
         FrameMap frameMap = new FrameMap();
         JetTypeMapper typeMapper = new JetTypeMapper(JetStandardLibrary.getJetStandardLibrary(namespace.getProject()), bindingContext);
-        ExpressionCodegen codegen = new ExpressionCodegen(mv, bindingContext, frameMap, typeMapper, Type.VOID_TYPE, null, OwnerKind.NAMESPACE);
+        ExpressionCodegen codegen = new ExpressionCodegen(mv, bindingContext, frameMap, typeMapper, null, Type.VOID_TYPE, null, OwnerKind.NAMESPACE);
 
         for (JetDeclaration declaration : namespace.getDeclarations()) {
             if (declaration instanceof JetProperty) {

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

@@ -96,7 +96,7 @@ public class PropertyCodegen {
         if (getter != null) {
             if (getter.getBodyExpression() != null) {
                 functionCodegen.generateMethod(getter, kind, mapper.mapGetterSignature(propertyDescriptor),
-                        Collections.<ValueParameterDescriptor>emptyList());
+                        null, Collections.<ValueParameterDescriptor>emptyList());
             }
             else if (!getter.hasModifier(JetTokens.PRIVATE_KEYWORD)) {
                 generateDefaultGetter(p, getter, kind);
@@ -118,7 +118,7 @@ public class PropertyCodegen {
                 final PropertySetterDescriptor setterDescriptor = propertyDescriptor.getSetter();
                 assert setterDescriptor != null;
                 functionCodegen.generateMethod(setter, kind, mapper.mapSetterSignature(propertyDescriptor),
-                        setterDescriptor.getUnsubstitutedValueParameters());
+                        null, setterDescriptor.getUnsubstitutedValueParameters());
             }
             else if (!p.hasModifier(JetTokens.PRIVATE_KEYWORD)) {
                 generateDefaultSetter(p, setter, kind);

idea/testData/codegen/extensionFunctions/simple.jet

@@ -0,0 +1,3 @@
+fun StringBuilder.first() = this.charAt(0)
+
+fun foo() = new StringBuilder("foo").first()

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

@@ -0,0 +1,20 @@
+package org.jetbrains.jet.codegen;
+
+import java.lang.reflect.Method;
+
+/**
+ * @author yole
+ */
+public class ExtensionFunctionsTest extends CodegenTestCase {
+    @Override
+    protected String getPrefix() {
+        return "extensionFunctions";
+    }
+
+    public void testSimple() throws Exception {
+        loadFile();
+        final Method foo = generateFunction("foo");
+        final Character c = (Character) foo.invoke(null);
+        assertEquals('f', c.charValue());
+    }
+}