Builtin iterator methods without the default library: a.map(f).map(g)

tslang/include/TypeScript/MLIRLogic/MLIRCodeLogic.h

@@ -1771,37 +1771,7 @@ class MLIRPropertyAccessCodeLogic
                     elementType = arrayType.getElementType();
                 }
 
-                auto isReduce = propName == "reduce";
-                SmallVector<mlir::Type> resultArgs;
-                if (isArrayCustomMethodReturnsBool(propName))
-                {
-                    resultArgs.push_back(mlir_ts::BooleanType::get(builder.getContext()));
-                }
-
-                mlir::Type genericTypeT;
-                SmallVector<mlir::Type> lambdaArgs{elementType};
-                if (isReduce)
-                {
-                    // add sum param
-                    genericTypeT = mlir_ts::NamedGenericType::get(builder.getContext(), mlir::FlatSymbolRefAttr::get(builder.getContext(), "T"));
-                    lambdaArgs.insert(&lambdaArgs.front(), genericTypeT);
-                }
-
-                auto lambdaFuncType = mlir_ts::FunctionType::get(builder.getContext(), lambdaArgs, resultArgs);
-
-                SmallVector<mlir::Type> funcArgs{lambdaFuncType};
-                if (isReduce)
-                {
-                    funcArgs.push_back(genericTypeT);
-                }
-
-                auto funcType = mlir_ts::FunctionType::get(builder.getContext(), funcArgs, resultArgs);
-                auto symbOp = builder.create<mlir_ts::ThisSymbolRefOp>(
-                    location, funcType, expression,
-                    mlir::FlatSymbolRefAttr::get(builder.getContext(), 
-                    getArrayCustomMethodName(propName)));
-                symbOp->setAttr(BUILTIN_FUNC_ATTR_NAME, mlir::BoolAttr::get(builder.getContext(), true));
-                return symbOp;
+                return CustomMethod(elementType);
             }
 
             return mlir::Value();
@@ -1810,6 +1780,51 @@ class MLIRPropertyAccessCodeLogic
         return mlir::Value();
     }
 
+    // The custom methods an iterator has too: `reduce`'s builtin takes an array (`T[]`).
+    bool isIteratorCustomMethod()
+    {
+        auto propName = getName();
+        return isArrayCustomMethod(propName) && propName != "reduce";
+    }
+
+    // The builtin custom method `getName()` of `expression`, whose elements are `elementType`:
+    // an array, or an iterator, which the builtin goes through with for...of.
+    mlir::Value CustomMethod(mlir::Type elementType)
+    {
+        auto propName = getName();
+        auto isReduce = propName == "reduce";
+        SmallVector<mlir::Type> resultArgs;
+        if (isArrayCustomMethodReturnsBool(propName))
+        {
+            resultArgs.push_back(mlir_ts::BooleanType::get(builder.getContext()));
+        }
+
+        mlir::Type genericTypeT;
+        SmallVector<mlir::Type> lambdaArgs{elementType};
+        if (isReduce)
+        {
+            // add sum param
+            genericTypeT = mlir_ts::NamedGenericType::get(builder.getContext(), mlir::FlatSymbolRefAttr::get(builder.getContext(), "T"));
+            lambdaArgs.insert(&lambdaArgs.front(), genericTypeT);
+        }
+
+        auto lambdaFuncType = mlir_ts::FunctionType::get(builder.getContext(), lambdaArgs, resultArgs);
+
+        SmallVector<mlir::Type> funcArgs{lambdaFuncType};
+        if (isReduce)
+        {
+            funcArgs.push_back(genericTypeT);
+        }
+
+        auto funcType = mlir_ts::FunctionType::get(builder.getContext(), funcArgs, resultArgs);
+        auto symbOp = builder.create<mlir_ts::ThisSymbolRefOp>(
+            location, funcType, expression,
+            mlir::FlatSymbolRefAttr::get(builder.getContext(), 
+            getArrayCustomMethodName(propName)));
+        symbOp->setAttr(BUILTIN_FUNC_ATTR_NAME, mlir::BoolAttr::get(builder.getContext(), true));
+        return symbOp;
+    }
+
     template <typename T> mlir::Value Ref(T refType)
     {
         if (auto constTupleType = dyn_cast<mlir_ts::ConstTupleType>(refType.getElementType()))

tslang/lib/TypeScript/MLIRGenAccessCall.cpp

@@ -385,6 +385,19 @@ namespace mlirgen
             }
         }
 
+        // Without the default library, an array's map/filter/forEach/every/some are builtins, and
+        // what map and filter give is a generator. An iterator has them too, as the default
+        // library's __Iterator gives them (`a.map(f).map(g)`, `g().filter(f)`): each builtin is a
+        // for...of over its receiver. After extension functions, so a user's own `map` wins.
+        if (!value && compileOptions.enableBuiltins && compileOptions.noDefaultLib && cl.isIteratorCustomMethod())
+        {
+            if (auto elementType = getIteratorValueType(location, objectValue, genContext))
+            {
+                cl.setExpression(objectValue);
+                return cl.CustomMethod(elementType);
+            }
+        }
+
         if (!value)
         {
             // During a speculative discovery/dummy run (e.g. inferring an enclosing

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -7756,8 +7756,9 @@ class MLIRGenImpl
             nf.createForOfStatement(undefined, declList, _src_array_ident,
                                     nf.createExpressionStatement(_yield_expr));
 
-        // iterator
-        auto iterName = MLIRHelper::getAnonymousName(location, ".iter", getFullNamespaceName());
+        // iterator; named by where its callback is too: `a.map(f).map(g)`'s two calls start together
+        auto iterName = MLIRHelper::getAnonymousName(
+            mlir::FusedLoc::get(builder.getContext(), {location, funcSrc.getLoc()}), ".iter", getFullNamespaceName());
 
         NodeArray<Statement> statements;
         statements.push_back(forOfStat);
@@ -7817,8 +7818,9 @@ class MLIRGenImpl
                                  nf.createExpressionStatement(_yield_expr),
                                  undefined));
 
-        // iterator
-        auto iterName = MLIRHelper::getAnonymousName(location, ".iter", getFullNamespaceName());
+        // iterator; named by where its callback is too: `a.map(f).map(g)`'s two calls start together
+        auto iterName = MLIRHelper::getAnonymousName(
+            mlir::FusedLoc::get(builder.getContext(), {location, funcSrc.getLoc()}), ".iter", getFullNamespaceName());
 
         NodeArray<Statement> statements;
         statements.push_back(forOfStat);
@@ -8217,6 +8219,12 @@ class MLIRGenImpl
         OperandsProcessingInfo* operandsProcessingInfo, const GenContext &genContext);
 
     bool hasIterator(mlir::Location location, mlir::Value source, const GenContext &genContext) 
+    {
+        return !!getIteratorValueType(location, source, genContext);
+    }
+
+    // The type of `value` in what `source.next()` returns, or null when `source` is not an iterator.
+    mlir::Type getIteratorValueType(mlir::Location location, mlir::Value source, const GenContext &genContext) 
     {
         auto nextPropertyType = evaluateProperty(location, source, ITERATOR_NEXT, genContext);
         if (nextPropertyType)
@@ -8243,14 +8251,15 @@ class MLIRGenImpl
                     });
 
                 auto propValue = mlir::StringAttr::get(builder.getContext(), "value");
-                if (std::any_of(fields.begin(), fields.end(), [&] (auto field) { return field.id == propValue; }))
+                auto valueField = std::find_if(fields.begin(), fields.end(), [&] (auto field) { return field.id == propValue; });
+                if (valueField != fields.end())
                 {
-                    return true;
+                    return valueField->type;
                 }
             }
         }
 
-        return false;
+        return mlir::Type();
     }
 
     bool isArrayLike(mlir::Location location, mlir::Value source, const GenContext &genContext) 

tslang/test/tester/CMakeLists.txt

@@ -619,6 +619,7 @@ tslang_add_test(NAME test-compile-02-iterator COMMAND test-runner "${PROJECT_SOU
 tslang_add_test(NAME test-compile-03-iterator COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/03iterator.ts")
 tslang_add_test(NAME test-compile-00-iterator-bug COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00iterator_bug.ts")
 tslang_add_test(NAME test-compile-00-map COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00map.ts")
+tslang_add_test(NAME test-compile-00-iterator-methods COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00iterator_methods.ts")
 tslang_add_test(NAME test-compile-01-map COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/01map.ts")
 tslang_add_test(NAME test-compile-00-filter COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00filter.ts")
 tslang_add_test(NAME test-compile-00-reduce COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00reduce.ts")
@@ -1105,6 +1106,7 @@ tslang_add_test(NAME test-jit-02-iterator COMMAND test-runner -jit "${PROJECT_SO
 tslang_add_test(NAME test-jit-03-iterator COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/03iterator.ts")
 tslang_add_test(NAME test-jit-00-iterator-bug COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00iterator_bug.ts")
 tslang_add_test(NAME test-jit-00-map COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00map.ts")
+tslang_add_test(NAME test-jit-00-iterator-methods COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00iterator_methods.ts")
 tslang_add_test(NAME test-jit-01-map COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/01map.ts")
 tslang_add_test(NAME test-jit-00-filter COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00filter.ts")
 tslang_add_test(NAME test-jit-00-reduce COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00reduce.ts")
@@ -1785,6 +1787,7 @@ set(TSLANG_CORPUS
     00intersection_type_generic.ts
     00intersection_type.ts
     00iterator_bug.ts
+    00iterator_methods.ts
     00iterator.ts
     00lambdas_generic_global.ts
     00lambdas.ts

tslang/test/tester/tests/00iterator_methods.ts

@@ -0,0 +1,46 @@
+// What the builtin map and filter give is a generator, and an iterator has the same builtin
+// methods as an array: map, filter, forEach, every, some.
+function* count(n: number) {
+    for (let i = 1; i <= n; i++) yield i;
+}
+
+function main() {
+    const a: number[] = [1, 2, 3, 4];
+
+    // a chain in one expression: both calls start at the same place
+    let s = 0;
+    for (const v of a.map(x => x * 2).map(x => x + 1)) s += v;
+    assert(s == 24);
+
+    s = 0;
+    for (const v of a.map(x => x * 2).filter(x => x > 2).map(x => x + 1)) s += v;
+    assert(s == 21);
+
+    // through a variable
+    const m = a.filter(x => x > 1);
+    s = 0;
+    m.forEach(x => { s += x; });
+    assert(s == 9);
+
+    assert(a.map(x => x * 2).every(x => x % 2 == 0));
+    assert(!a.map(x => x * 2).every(x => x > 2));
+    assert(a.map(x => x * 2).some(x => x > 7));
+    assert(!a.filter(x => x > 2).some(x => x > 4));
+
+    // a generator of its own
+    s = 0;
+    for (const v of count(5).filter(x => x % 2 == 1).map(x => x * 10)) s += v;
+    assert(s == 90);
+
+    // the result is still an iterator
+    const it = a.map(x => x + 100);
+    const first = it.next();
+    assert(!first.done && first.value == 101);
+
+    const strs = ["a", "b"];
+    let joined = "";
+    strs.map(x => x + "_").forEach(x => { joined += x; });
+    assert(joined == "a_b_");
+
+    print("done.");
+}