diff --git a/tslang/include/TypeScript/MLIRLogic/MLIRCodeLogic.h b/tslang/include/TypeScript/MLIRLogic/MLIRCodeLogic.h index da48336c3..7ffe9e6ff 100644 --- a/tslang/include/TypeScript/MLIRLogic/MLIRCodeLogic.h +++ b/tslang/include/TypeScript/MLIRLogic/MLIRCodeLogic.h @@ -1771,37 +1771,7 @@ class MLIRPropertyAccessCodeLogic elementType = arrayType.getElementType(); } - auto isReduce = propName == "reduce"; - SmallVector resultArgs; - if (isArrayCustomMethodReturnsBool(propName)) - { - resultArgs.push_back(mlir_ts::BooleanType::get(builder.getContext())); - } - - mlir::Type genericTypeT; - SmallVector 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 funcArgs{lambdaFuncType}; - if (isReduce) - { - funcArgs.push_back(genericTypeT); - } - - auto funcType = mlir_ts::FunctionType::get(builder.getContext(), funcArgs, resultArgs); - auto symbOp = builder.create( - 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 resultArgs; + if (isArrayCustomMethodReturnsBool(propName)) + { + resultArgs.push_back(mlir_ts::BooleanType::get(builder.getContext())); + } + + mlir::Type genericTypeT; + SmallVector 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 funcArgs{lambdaFuncType}; + if (isReduce) + { + funcArgs.push_back(genericTypeT); + } + + auto funcType = mlir_ts::FunctionType::get(builder.getContext(), funcArgs, resultArgs); + auto symbOp = builder.create( + 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 mlir::Value Ref(T refType) { if (auto constTupleType = dyn_cast(refType.getElementType())) diff --git a/tslang/lib/TypeScript/MLIRGenAccessCall.cpp b/tslang/lib/TypeScript/MLIRGenAccessCall.cpp index 7e0740c65..16f4c69b1 100644 --- a/tslang/lib/TypeScript/MLIRGenAccessCall.cpp +++ b/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 diff --git a/tslang/lib/TypeScript/MLIRGenImpl.h b/tslang/lib/TypeScript/MLIRGenImpl.h index 245d63b85..6033734bd 100644 --- a/tslang/lib/TypeScript/MLIRGenImpl.h +++ b/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 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 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) diff --git a/tslang/test/tester/CMakeLists.txt b/tslang/test/tester/CMakeLists.txt index 1bebfe032..e1e793200 100644 --- a/tslang/test/tester/CMakeLists.txt +++ b/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 diff --git a/tslang/test/tester/tests/00iterator_methods.ts b/tslang/test/tester/tests/00iterator_methods.ts new file mode 100644 index 000000000..d674ec317 --- /dev/null +++ b/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."); +}