JS backend: fixed variable capturing in extension functions.

Fixed ClosureTest.testWrappedVariableInExtensionFun IteratorsTest.foldReducesTheFirstNElements and IteratorsTest.takeExtractsTheFirstNElements. All tests passed.

(cherry picked from commit f8f1895)

js/js.tests/test/org/jetbrains/k2js/test/semantics/ClosureTest.java

@@ -47,4 +47,8 @@ public final class ClosureTest extends SingleFileTranslationTest {
     public void testClosureInNestedFunctionsInMethod() throws Exception {
         checkFooBoxIsOk();
     }
+
+    public void testWrappedVariableInExtensionFun() throws Exception {
+        fooBoxTest();
+    }
 }

js/js.translator/src/org/jetbrains/k2js/translate/context/AliasingContext.java

@@ -18,19 +18,19 @@ package org.jetbrains.k2js.translate.context;
 
 import com.google.dart.compiler.backend.js.ast.JsExpression;
 import com.google.dart.compiler.backend.js.ast.JsName;
+import gnu.trove.THashMap;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
 import org.jetbrains.jet.lang.psi.JetExpression;
 
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Map;
 
 public class AliasingContext {
     private static final AliasingContext ROOT = new AliasingContext(null) {
         @Override
-        public JsExpression getAliasForDescriptor(@NotNull DeclarationDescriptor descriptor) {
+        protected JsExpression getAliasForDescriptor(@NotNull DeclarationDescriptor descriptor, boolean fromChild) {
             return null;
         }
 
@@ -38,6 +38,11 @@ public class AliasingContext {
         public JsName getAliasForExpression(@NotNull JetExpression element) {
             return null;
         }
+
+        @Override
+        public void registerAlias(@NotNull DeclarationDescriptor descriptor, @NotNull JsExpression alias) {
+            throw new IllegalStateException();
+        }
     };
 
     public static AliasingContext getCleanContext() {
@@ -53,7 +58,7 @@ public class AliasingContext {
     @Nullable
     private final AliasingContext parent;
 
-    private AliasingContext(@Nullable AliasingContext parent) {
+    /*package*/ AliasingContext(@Nullable AliasingContext parent) {
         this(parent, null, null);
     }
 
@@ -72,6 +77,20 @@ public class AliasingContext {
         return new AliasingContext(this, Collections.singletonMap(descriptor, alias), null);
     }
 
+    @NotNull
+    public AliasingContext notShareableThisAliased(@NotNull final DeclarationDescriptor thisDescriptor, @NotNull final JsExpression alias) {
+        return new AliasingContext(this) {
+            @Nullable
+            @Override
+            protected JsExpression getAliasForDescriptor(@NotNull DeclarationDescriptor descriptor, boolean fromChild) {
+                if (!fromChild && thisDescriptor == descriptor) {
+                    return alias;
+                }
+                return super.getAliasForDescriptor(descriptor, fromChild);
+            }
+        };
+    }
+
     @NotNull
     public AliasingContext withExpressionsAliased(@NotNull Map<JetExpression, JsName> aliasesForExpressions) {
         return new AliasingContext(this, null, aliasesForExpressions);
@@ -82,22 +101,17 @@ public class AliasingContext {
         return new AliasingContext(this, aliases, null);
     }
 
+
     @Nullable
     public JsExpression getAliasForDescriptor(@NotNull DeclarationDescriptor descriptor) {
-        if (aliasesForDescriptors == null) {
+        // these aliases cannot be shared and applicable only in current context
-            return null;
+        return getAliasForDescriptor(descriptor, false);
-        }
+    }
 
-        JsExpression alias = aliasesForDescriptors.get(descriptor.getOriginal());
+    @Nullable
-        if (alias != null) {
+    protected JsExpression getAliasForDescriptor(@NotNull DeclarationDescriptor descriptor, boolean fromChild) {
-            return alias;
+        JsExpression alias = aliasesForDescriptors == null ? null : aliasesForDescriptors.get(descriptor.getOriginal());
-        }
+        return alias != null || parent == null ? alias : parent.getAliasForDescriptor(descriptor, true);
-
-        if (parent == null) {
-            return null;
-        }
-
-        return parent.getAliasForDescriptor(descriptor);
     }
 
     @Nullable
@@ -106,6 +120,13 @@ public class AliasingContext {
         return alias != null || parent == null ? alias : parent.getAliasForExpression(element);
     }
 
+    /**
+     * Usages:
+     * 1) Local variable captured in closure. If captured in closure, any modification in closure should affect captured variable.
+     * So, "var count = n" wrapped as "var count = {v: n}". descriptor wil be property descriptor, alias will be JsObjectLiteral
+     *
+     * 2) Local named function.
+     */
     public void registerAlias(@NotNull DeclarationDescriptor descriptor, @NotNull JsExpression alias) {
         if (aliasesForDescriptors == null) {
             aliasesForDescriptors = Collections.singletonMap(descriptor, alias);
@@ -113,7 +134,7 @@ public class AliasingContext {
         else {
             if (aliasesForDescriptors.size() == 1) {
                 Map<DeclarationDescriptor, JsExpression> singletonMap = aliasesForDescriptors;
-                aliasesForDescriptors = new HashMap<DeclarationDescriptor, JsExpression>();
+                aliasesForDescriptors = new THashMap<DeclarationDescriptor, JsExpression>();
                 aliasesForDescriptors.put(singletonMap.keySet().iterator().next(), singletonMap.values().iterator().next());
             }
             JsExpression prev = aliasesForDescriptors.put(descriptor, alias);

js/js.translator/src/org/jetbrains/k2js/translate/context/Namer.java

@@ -32,6 +32,7 @@ public final class Namer {
     public static final String KOTLIN_LOWER_NAME = KOTLIN_NAME.toLowerCase();
 
     public static final String CALLEE_NAME = "$fun";
+    public static final String OUTER_CLASS_NAME = "$outer";
 
     private static final String INITIALIZE_METHOD_NAME = "initialize";
     private static final String CLASS_OBJECT_NAME = "createClass";
@@ -42,7 +43,7 @@ public final class Namer {
     private static final String BACKING_FIELD_PREFIX = "$";
     private static final String SUPER_METHOD_NAME = "super_init";
     private static final String ROOT_NAMESPACE = "_";
-    private static final String RECEIVER_PARAMETER_NAME = "receiver";
+    private static final String RECEIVER_PARAMETER_NAME = "$receiver";
     private static final String CLASSES_OBJECT_NAME = "classes";
     private static final String THROW_NPE_FUN_NAME = "throwNPE";
 

js/js.translator/src/org/jetbrains/k2js/translate/context/TranslationContext.java

@@ -61,10 +61,12 @@ public class TranslationContext {
         return staticContext.isEcma5();
     }
 
-    private TranslationContext(@NotNull StaticContext staticContext,
+    private TranslationContext(
+            @NotNull StaticContext staticContext,
             @NotNull DynamicContext dynamicContext,
             @NotNull AliasingContext aliasingContext,
-            @Nullable UsageTracker usageTracker) {
+            @Nullable UsageTracker usageTracker
+    ) {
         this.dynamicContext = dynamicContext;
         this.staticContext = staticContext;
         this.aliasingContext = aliasingContext;
@@ -72,10 +74,17 @@ public class TranslationContext {
     }
 
     private TranslationContext(@NotNull TranslationContext parent, @NotNull AliasingContext aliasingContext) {
-        dynamicContext = parent.dynamicContext;
+        this(parent.staticContext, parent.dynamicContext, aliasingContext, parent.usageTracker);
-        staticContext = parent.staticContext;
+    }
-        this.aliasingContext = aliasingContext;
+
-        usageTracker = parent.usageTracker;
+    private TranslationContext(
+            @NotNull TranslationContext parent,
+            @NotNull JsFunction fun,
+            @NotNull AliasingContext aliasingContext,
+            @Nullable UsageTracker usageTracker
+    ) {
+        this(parent.staticContext, DynamicContext.newContext(fun.getScope(), fun.getBody()), aliasingContext,
+             usageTracker == null ? parent.usageTracker : usageTracker);
     }
 
     public UsageTracker usageTracker() {
@@ -88,7 +97,7 @@ public class TranslationContext {
 
     @NotNull
     public TranslationContext contextWithScope(@NotNull JsFunction fun) {
-        return contextWithScope(fun, aliasingContext, usageTracker);
+        return new TranslationContext(this, fun, aliasingContext, null);
     }
 
     @NotNull
@@ -100,8 +109,13 @@ public class TranslationContext {
     }
 
     @NotNull
-    public TranslationContext contextWithScope(@NotNull JsFunction fun, @NotNull AliasingContext aliasingContext, @Nullable UsageTracker usageTracker) {
+    public TranslationContext newFunctionBody(
-        return contextWithScope(fun.getScope(), fun.getBody(), aliasingContext, usageTracker);
+            @NotNull JsFunction fun,
+            @Nullable AliasingContext aliasingContext,
+            @Nullable UsageTracker usageTracker
+    ) {
+        return new TranslationContext(this, fun, aliasingContext == null ? new AliasingContext(this.aliasingContext) : aliasingContext,
+                                      usageTracker);
     }
 
     @NotNull
@@ -115,8 +129,8 @@ public class TranslationContext {
     }
 
     @NotNull
-    public TranslationContext innerContextWithThisAliased(@NotNull DeclarationDescriptor correspondingDescriptor, @NotNull JsName alias) {
+    public TranslationContext innerContextWithThisAliased(@NotNull DeclarationDescriptor correspondingDescriptor, @NotNull JsNameRef alias) {
-        return new TranslationContext(this, aliasingContext.inner(correspondingDescriptor, alias.makeRef()));
+        return new TranslationContext(this, aliasingContext.inner(correspondingDescriptor, alias));
     }
 
     @NotNull

js/js.translator/src/org/jetbrains/k2js/translate/declaration/DeclarationBodyVisitor.java

@@ -57,8 +57,7 @@ public class DeclarationBodyVisitor extends TranslatorVisitor<Void> {
     }
 
     @Override
-    public Void visitNamedFunction(@NotNull JetNamedFunction expression,
+    public Void visitNamedFunction(@NotNull JetNamedFunction expression, @NotNull TranslationContext context) {
-            @NotNull TranslationContext context) {
         FunctionDescriptor descriptor = getFunctionDescriptor(context.bindingContext(), expression);
         if (descriptor.getModality() == Modality.ABSTRACT) {
             return null;

js/js.translator/src/org/jetbrains/k2js/translate/expression/ExpressionVisitor.java

@@ -39,10 +39,10 @@ import org.jetbrains.k2js.translate.operation.BinaryOperationTranslator;
 import org.jetbrains.k2js.translate.operation.UnaryOperationTranslator;
 import org.jetbrains.k2js.translate.reference.*;
 import org.jetbrains.k2js.translate.utils.BindingUtils;
+import org.jetbrains.k2js.translate.utils.JsAstUtils;
 import org.jetbrains.k2js.translate.utils.TranslationUtils;
 import org.jetbrains.k2js.translate.utils.mutator.AssignToExpressionMutator;
 
-import java.util.Collections;
 import java.util.List;
 
 import static org.jetbrains.k2js.translate.general.Translation.translateAsExpression;
@@ -149,15 +149,13 @@ public final class ExpressionVisitor extends TranslatorVisitor<JsNode> {
     @NotNull
     // assume it is a local variable declaration
     public JsNode visitProperty(@NotNull JetProperty expression, @NotNull TranslationContext context) {
-        VariableDescriptor descriptor = context.bindingContext().get(BindingContext.VARIABLE, expression);
+        VariableDescriptor descriptor = BindingContextUtils.getNotNull(context.bindingContext(), BindingContext.VARIABLE, expression);
-        assert descriptor != null;
         JsExpression initializer = translateInitializerForProperty(expression, context);
         JsName name = context.getNameForDescriptor(descriptor);
         if (descriptor.isVar() && context.bindingContext().get(BindingContext.CAPTURED_IN_CLOSURE, descriptor) != null) {
             // well, wrap it
             JsNameRef alias = new JsNameRef("v", new JsNameRef(name));
-            initializer = new JsObjectLiteral(
+            initializer = JsAstUtils.wrapValue(alias, initializer == null ? JsLiteral.NULL : initializer);
-                    Collections.singletonList(new JsPropertyInitializer(alias, initializer == null ? JsLiteral.NULL : initializer)));
             context.aliasingContext().registerAlias(descriptor, alias);
         }
 

js/js.translator/src/org/jetbrains/k2js/translate/expression/FunctionTranslator.java

@@ -21,6 +21,7 @@ import com.google.dart.compiler.backend.js.ast.JsFunction;
 import com.google.dart.compiler.backend.js.ast.JsName;
 import com.google.dart.compiler.backend.js.ast.JsParameter;
 import com.google.dart.compiler.backend.js.ast.JsPropertyInitializer;
+import com.intellij.util.SmartList;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
@@ -29,14 +30,15 @@ import org.jetbrains.jet.lang.descriptors.Modality;
 import org.jetbrains.jet.lang.descriptors.ValueParameterDescriptor;
 import org.jetbrains.jet.lang.psi.JetDeclarationWithBody;
 import org.jetbrains.jet.lang.psi.JetExpression;
-import org.jetbrains.jet.lang.psi.JetFunctionLiteral;
+import org.jetbrains.jet.lang.psi.JetFunctionLiteralExpression;
+import org.jetbrains.k2js.translate.context.AliasingContext;
 import org.jetbrains.k2js.translate.context.Namer;
 import org.jetbrains.k2js.translate.context.TranslationContext;
 import org.jetbrains.k2js.translate.general.AbstractTranslator;
 import org.jetbrains.k2js.translate.utils.JsDescriptorUtils;
 import org.jetbrains.k2js.translate.utils.TranslationUtils;
 
-import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 import static org.jetbrains.k2js.translate.utils.BindingUtils.getFunctionDescriptor;
@@ -76,24 +78,18 @@ public final class FunctionTranslator extends AbstractTranslator {
 
     @NotNull
     private TranslationContext getFunctionBodyContext() {
+        AliasingContext aliasingContext;
         if (isExtensionFunction()) {
-            return getFunctionBodyContextForExtensionFunction();
+            DeclarationDescriptor expectedReceiverDescriptor = getExpectedReceiverDescriptor(descriptor);
+            assert expectedReceiverDescriptor != null;
+            extensionFunctionReceiverName = functionObject.getScope().declareName(Namer.getReceiverParameterName());
+            //noinspection ConstantConditions
+            aliasingContext = context().aliasingContext().inner(expectedReceiverDescriptor, extensionFunctionReceiverName.makeRef());
         }
-        return getContextWithFunctionBodyBlock();
+        else {
-    }
+            aliasingContext = null;
-
+        }
-    @NotNull
+        return context().newFunctionBody(functionObject, aliasingContext, null);
-    private TranslationContext getFunctionBodyContextForExtensionFunction() {
-        TranslationContext contextWithFunctionBodyBlock = getContextWithFunctionBodyBlock();
-        extensionFunctionReceiverName = contextWithFunctionBodyBlock.scope().declareName(Namer.getReceiverParameterName());
-        DeclarationDescriptor expectedReceiverDescriptor = getExpectedReceiverDescriptor(descriptor);
-        assert expectedReceiverDescriptor != null;
-        return contextWithFunctionBodyBlock.innerContextWithThisAliased(expectedReceiverDescriptor, extensionFunctionReceiverName);
-    }
-
-    @NotNull
-    private TranslationContext getContextWithFunctionBodyBlock() {
-        return context().contextWithScope(functionObject);
     }
 
     @NotNull
@@ -125,7 +121,11 @@ public final class FunctionTranslator extends AbstractTranslator {
 
     @NotNull
     private List<JsParameter> translateParameters() {
-        List<JsParameter> jsParameters = new ArrayList<JsParameter>();
+        if (extensionFunctionReceiverName == null && descriptor.getValueParameters().isEmpty()) {
+            return Collections.emptyList();
+        }
+
+        List<JsParameter> jsParameters = new SmartList<JsParameter>();
         mayBeAddThisParameterForExtensionFunction(jsParameters);
         addParameters(jsParameters, descriptor, context());
         return jsParameters;
@@ -145,6 +145,6 @@ public final class FunctionTranslator extends AbstractTranslator {
     }
 
     private boolean isExtensionFunction() {
-        return JsDescriptorUtils.isExtension(descriptor) && !(functionDeclaration instanceof JetFunctionLiteral);
+        return JsDescriptorUtils.isExtension(descriptor) && !(functionDeclaration instanceof JetFunctionLiteralExpression);
     }
 }

js/js.translator/src/org/jetbrains/k2js/translate/expression/LiteralFunctionTranslator.java

@@ -76,15 +76,16 @@ public class LiteralFunctionTranslator {
         TranslationContext funContext;
         boolean asInner;
         ClassDescriptor outerClass;
-        AliasingContext aliasingContext = rootContext.aliasingContext();
+        AliasingContext aliasingContext;
         DeclarationDescriptor receiverDescriptor = getExpectedReceiverDescriptor(descriptor);
         JsName receiverName;
         if (receiverDescriptor == null) {
             receiverName = null;
+            aliasingContext = null;
         }
         else {
             receiverName = fun.getScope().declareName(Namer.getReceiverParameterName());
-            aliasingContext = aliasingContext.inner(receiverDescriptor, receiverName.makeRef());
+            aliasingContext = outerContext.aliasingContext().inner(receiverDescriptor, receiverName.makeRef());
         }
 
         if (descriptor.getContainingDeclaration() instanceof ConstructorDescriptor) {
@@ -95,7 +96,7 @@ public class LiteralFunctionTranslator {
             assert outerClass != null;
 
             if (receiverDescriptor == null) {
-                aliasingContext = aliasingContext.inner(outerClass, new JsNameRef("o", fun.getName().makeRef()));
+                aliasingContext = outerContext.aliasingContext().notShareableThisAliased(outerClass, new JsNameRef("o", fun.getName().makeRef()));
             }
         }
         else {
@@ -103,8 +104,8 @@ public class LiteralFunctionTranslator {
             asInner = DescriptorUtils.isTopLevelDeclaration(descriptor);
         }
 
-        funContext = rootContext.contextWithScope(fun, aliasingContext,
+        funContext = outerContext.newFunctionBody(fun, aliasingContext,
-                                                  new UsageTracker(descriptor, outerContext.usageTracker(), outerClass));
+                                                 new UsageTracker(descriptor, outerContext.usageTracker(), outerClass));
 
         fun.getBody().getStatements().addAll(translateFunctionBody(descriptor, declaration, funContext).getStatements());
 
@@ -160,10 +161,10 @@ public class LiteralFunctionTranslator {
             @NotNull ClassDescriptor descriptor,
             @NotNull ClassTranslator classTranslator) {
         JsFunction fun = createFunction();
-        JsNameRef outerClassRef = fun.getScope().declareName("$this").makeRef();
+        JsNameRef outerClassRef = fun.getScope().declareName(Namer.OUTER_CLASS_NAME).makeRef();
         UsageTracker usageTracker = new UsageTracker(descriptor, null, outerClass);
-        TranslationContext funContext = rootContext.contextWithScope(fun, rootContext.aliasingContext().inner(outerClass, outerClassRef),
+        TranslationContext funContext = rootContext.newFunctionBody(fun, rootContext.aliasingContext().inner(outerClass, outerClassRef),
-                                                                     usageTracker);
+                                                                    usageTracker);
 
         fun.getBody().getStatements().add(new JsReturn(classTranslator.translate(funContext)));
         JetClassBody body = declaration.getBody();

js/js.translator/src/org/jetbrains/k2js/translate/reference/InlinedCallExpressionTranslator.java

@@ -129,11 +129,10 @@ public final class InlinedCallExpressionTranslator extends AbstractCallExpressio
     @NotNull
     private TranslationContext contextWithAlias(@NotNull TranslationContext contextWithAliasForThisExpression,
                                                 @NotNull JsExpression aliasExpression, @Nullable DeclarationDescriptor descriptorToAlias) {
-        TranslationContext newContext = contextWithAliasForThisExpression;
         TemporaryVariable aliasForReceiver = context().declareTemporary(aliasExpression);
         assert descriptorToAlias != null;
-        newContext =
+        TranslationContext newContext =
-            newContext.innerContextWithThisAliased(descriptorToAlias, aliasForReceiver.name());
+                contextWithAliasForThisExpression.innerContextWithThisAliased(descriptorToAlias, aliasForReceiver.reference());
         newContext.addStatementToCurrentBlock(aliasForReceiver.assignmentExpression().makeStmt());
         return newContext;
     }

js/js.translator/src/org/jetbrains/k2js/translate/utils/JsAstUtils.java

@@ -281,4 +281,9 @@ public final class JsAstUtils {
         to.add(new JsInvocation(EMPTY_REF, new JsInvocation(packageBlockFunction)).makeStmt());
         return packageBlockFunction;
     }
+
+    @NotNull
+    public static JsObjectLiteral wrapValue(@NotNull JsExpression label, @NotNull JsExpression value) {
+        return new JsObjectLiteral(Collections.singletonList(new JsPropertyInitializer(label, value)));
+    }
 }

js/js.translator/testFiles/closure/cases/wrappedVariableInExtensionFun.kt

@@ -0,0 +1,10 @@
+package foo
+
+fun Any.foo(n: Int): ()->Boolean {
+    var count = n
+    return { --count >= 0 }
+}
+
+fun box(): Boolean {
+  return 1.foo(3)() && !1.foo(0)()
+}

libraries/stdlib/test/iterators/IteratorsTest.kt

@@ -16,14 +16,12 @@ class IteratorsTest {
         assertEquals(arrayList(144, 233, 377, 610, 987), fibonacci().filter { it > 100 }.takeWhile { it < 1000 }.toList())
     }
 
-    // TODO fix and enable this test
+    test fun foldReducesTheFirstNElements() {
-    fun foldReducesTheFirstNElements() {
         val sum = { (a: Int, b: Int) -> a + b }
         assertEquals(arrayList(13, 21, 34, 55, 89).fold(0, sum), fibonacci().filter { it > 10 }.take(5).fold(0, sum))
     }
 
-    // TODO fix and enable this test
+    test fun takeExtractsTheFirstNElements() {
-    fun takeExtractsTheFirstNElements() {
         assertEquals(arrayList(0, 1, 1, 2, 3, 5, 8, 13, 21, 34), fibonacci().take(10).toList())
     }