feat: pipeable index (#207)

* feat: pipeable index

* fix: typo

* refactor: combine checkTsPlusCustomCall and checkTsPlusCustomSignature

* refactor: cache only one signature

effect/packages/package1/src/prelude/pipeable/pipeable-index.ts

@@ -0,0 +1,23 @@
+/**
+ * @tsplus type pipeable-index/X
+ */
+export interface X {}
+
+// /**
+//  * @tsplus index pipeable-index/X
+//  */
+// export declare function f(self: X, index: number): void
+
+// /**
+//  * @tsplus pipeable-index pipeable-index/X
+//  */
+// export declare function fn(index: number): (self: X) => string
+
+/**
+ * @tsplus pipeable-index pipeable-index/X
+ */
+export declare const fn: (index: number) => (self: X) => string
+
+declare const x: X
+
+const y = x[1]

effect/packages/package1/src/type-driven-overloads.ts

@@ -8,7 +8,6 @@ T.make(1).op(1)
 
 T.make(1) + true
 
-// @ts-expect-error
 T.make(1).op(true)
 
 T.make(1)

src/compiler/checker.ts

@@ -379,8 +379,9 @@ namespace ts {
         const unresolvedStaticCache = new Map<string, ESMap<string, TsPlusUnresolvedStaticExtension>>();
         const identityCache = new Map<string, FunctionDeclaration>();
         const callCache = new Map<Node, TsPlusStaticFunctionExtension>();
-        const indexCache = new Map<string, { declaration: FunctionDeclaration, definition: SourceFile, exportName: string }>();
-        const indexAccessExpressionCache = new Map<Node, { declaration: FunctionDeclaration, definition: SourceFile, exportName: string }>();
+        const indexCache = new Map<string, () => { declaration: FunctionDeclaration | VariableDeclarationWithIdentifier, signature: Signature, definition: SourceFile, exportName: string } | undefined>();
+        const resolvedIndexCache = new Map<string, { declaration: FunctionDeclaration | VariableDeclarationWithIdentifier, signature: Signature, definition: SourceFile, exportName: string }>();
+        const indexAccessExpressionCache = new Map<Node, { signature: Signature, declaration: FunctionDeclaration | VariableDeclarationWithIdentifier, definition: SourceFile, exportName: string }>();
         const inheritanceSymbolCache = new Map<Symbol, Set<Symbol>>()
         const tsPlusGlobalImportCache = new Map<string, TsPlusGlobalImport>()
         const unificationInProgress = {
@@ -15495,11 +15496,15 @@ namespace ts {
             errorNode: Node,
             args: Expression[],
             checkMode: CheckMode | undefined,
+            signature?: Signature,
             addDiagnostic?: (_: Diagnostic) => void,
         ): Type {
             const funcType = getTypeOfNode(declaration);
             const apparentType = getApparentType(funcType);
-            const candidate = getSignaturesOfType(apparentType, SignatureKind.Call)[0]!;
+            const candidate = signature ?? getSignaturesOfType(apparentType, SignatureKind.Call)[0]!;
+            if (!candidate) {
+                return errorType;
+            }
             const node = factory.createCallExpression(
                 factory.createIdentifier("$tsplus_custom_call"),
                 [],
@@ -30931,7 +30936,7 @@ namespace ts {
             const indexer = tags.flatMap((tag) => {
                 const indexer = indexCache.get(tag)
                 if (indexer) {
-                    return [indexer]
+                    return [indexer()]
                 }
                 return []
             })[0]
@@ -30940,7 +30945,8 @@ namespace ts {
                     indexer.declaration,
                     node,
                     [node.expression, indexExpression],
-                    checkMode
+                    checkMode,
+                    indexer.signature,
                 )
                 if (!isErrorType(res)) {
                     indexAccessExpressionCache.set(node, indexer)
@@ -46087,6 +46093,12 @@ namespace ts {
             }
             return [];
         }
+        function collectTsPlusPipeableIndexTags(declaration: Declaration) {
+            if (isVariableDeclaration(declaration) || isFunctionDeclaration(declaration)) {
+                return declaration.tsPlusPipeableIndexTags || [];
+            }
+            return [];
+        }
         function collectTsPlusTypeTags(declaration: Declaration) {
             if (isInterfaceDeclaration(declaration) || isTypeAliasDeclaration(declaration) || isClassDeclaration(declaration)) {
                 return declaration.tsPlusTypeTags || [];
@@ -46414,7 +46426,7 @@ namespace ts {
             }
             return undefined;
         }
-        function isVairableDeclarationWithFunction(node: VariableDeclaration): node is VariableDeclarationWithFunction {
+        function isVariableDeclarationWithFunction(node: VariableDeclaration): node is VariableDeclarationWithFunction {
             return isIdentifier(node.name) && !!node.initializer && (isArrowFunction(node.initializer) || isFunctionExpression(node.initializer));
         }
         function isEachDefined<A>(array: (A | undefined)[]): array is A[] {
@@ -46431,7 +46443,7 @@ namespace ts {
             }
             const type = getTypeOfNode(pipeable);
             const signatures = getSignaturesOfType(type, SignatureKind.Call);
-            if (isVairableDeclarationWithFunction(pipeable)) {
+            if (isVariableDeclarationWithFunction(pipeable)) {
                 const initializer = pipeable.initializer;
                 const body = initializer.body;
                 let returnFn: ArrowFunction | FunctionTypeNode | undefined;
@@ -47124,15 +47136,72 @@ namespace ts {
             if(declaration.name && isIdentifier(declaration.name)) {
                 for (const target of collectTsPlusIndexTags(declaration)) {
                     tsPlusFiles.add(getSourceFileOfNode(declaration));
-                    tsPlusFiles.add(getSourceFileOfNode(declaration));
-                    indexCache.set(target, {
-                        declaration,
-                        definition: getSourceFileOfNode(declaration),
-                        exportName: declaration.name.escapedText.toString()
+                    indexCache.set(target, () => {
+                        if (resolvedIndexCache.has(target)) {
+                            return resolvedIndexCache.get(target);
+                        }
+                        const definition = getSourceFileOfNode(declaration);
+                        const signatures = getSignaturesOfType(getTypeOfNode(declaration), SignatureKind.Call) as Signature[];
+                        if (signatures[0]) {
+                            resolvedIndexCache.set(target, { declaration, signature: signatures[0], definition, exportName: declaration.name!.escapedText.toString() });
+                        }
+                        return resolvedIndexCache.get(target)!;
                     });
                 }
             }
         }
+        function cacheTsPlusIndexVariable(declaration: VariableDeclarationWithIdentifier) {
+            for (const target of collectTsPlusIndexTags(declaration)) {
+                tsPlusFiles.add(getSourceFileOfNode(declaration));
+                indexCache.set(target, () => {
+                    if (resolvedIndexCache.has(target)) {
+                        return resolvedIndexCache.get(target);
+                    }
+                    const definition = getSourceFileOfNode(declaration);
+                    const signatures = getSignaturesOfType(getTypeOfNode(declaration), SignatureKind.Call) as Signature[];
+                    if (signatures[0]) {
+                        resolvedIndexCache.set(target, { declaration, signature: signatures[0], definition, exportName: declaration.name!.escapedText.toString() });
+                    }
+                    return resolvedIndexCache.get(target)!;
+                });
+            }
+        }
+        function cacheTsPlusPipeableIndexFunction(declaration: FunctionDeclaration) {
+            if(declaration.name && isIdentifier(declaration.name)) {
+                for (const target of collectTsPlusPipeableIndexTags(declaration)) {
+                    tsPlusFiles.add(getSourceFileOfNode(declaration));
+                    indexCache.set(target, () => {
+                        if (resolvedIndexCache.has(target)) {
+                            return resolvedIndexCache.get(target);
+                        }
+                        const definition = getSourceFileOfNode(declaration);
+                        const exportName = declaration.name!.escapedText.toString();
+                        const typeAndSignatures = getTsPlusFluentSignatureForPipeableFunction(definition, exportName, exportName, declaration);
+                        if (typeAndSignatures && typeAndSignatures[1][0]) {
+                            resolvedIndexCache.set(target, { declaration, signature: typeAndSignatures[1][0], definition, exportName });
+                            return resolvedIndexCache.get(target);
+                        }
+                    })
+                }
+            }
+        }
+        function cacheTsPlusPipeableIndexVariable(declaration: VariableDeclarationWithIdentifier) {
+            for (const target of collectTsPlusPipeableIndexTags(declaration)) {
+                tsPlusFiles.add(getSourceFileOfNode(declaration));
+                indexCache.set(target, () => {
+                    if (resolvedIndexCache.has(target)) {
+                        return resolvedIndexCache.get(target);
+                    }
+                    const definition = getSourceFileOfNode(declaration);
+                    const exportName = declaration.name!.escapedText.toString();
+                    const typeAndSignatures = getTsPlusFluentSignatureForPipeableVariableDeclaration(definition, exportName, exportName, declaration as VariableDeclarationWithFunction | VariableDeclarationWithFunctionType);
+                    if (typeAndSignatures && typeAndSignatures[1][0]) {
+                        resolvedIndexCache.set(target, { declaration, signature: typeAndSignatures[1][0], definition, exportName });
+                        return resolvedIndexCache.get(target);
+                    }
+                })
+            }
+        }
         function collectTsPlusSymbols(file: SourceFile): void {
             for (const declaration of file.tsPlusContext.type) {
                 cacheTsPlusType(declaration);
@@ -47192,7 +47261,20 @@ namespace ts {
                 cacheTsPlusUnifyFunction(declaration);
             }
             for (const declaration of file.tsPlusContext.index) {
-                cacheTsPlusIndexFunction(declaration);                        
+                if (isFunctionDeclaration(declaration)) {
+                    cacheTsPlusIndexFunction(declaration);
+                }
+                else {
+                    cacheTsPlusIndexVariable(declaration);                        
+                }
+            }
+            for (const declaration of file.tsPlusContext.pipeableIndex) {
+                if (isFunctionDeclaration(declaration)) {
+                    cacheTsPlusPipeableIndexFunction(declaration);
+                }
+                else {
+                    cacheTsPlusPipeableIndexVariable(declaration);                        
+                }
             }
         }
         function fillTsPlusLocalScope(file: SourceFile) {
@@ -47261,6 +47343,7 @@ namespace ts {
             getterCache.clear();
             callCache.clear();
             indexCache.clear();
+            resolvedIndexCache.clear();
             indexAccessExpressionCache.clear();
             inheritanceSymbolCache.clear();
             tsPlusWorldScope.implicits.clear();

src/compiler/factory/nodeFactory.ts

@@ -5313,6 +5313,7 @@ namespace ts {
                 pipeable: [],
                 operator: [],
                 pipeableOperator: [],
+                pipeableIndex: [],
                 static: [],
                 getter: [],
                 unify: [],

src/compiler/parser.ts

@@ -1252,6 +1252,12 @@ namespace ts {
                             if (declaration.tsPlusGetterTags && declaration.tsPlusGetterTags.length > 0) {
                                 file.tsPlusContext.getter.push(declaration as VariableDeclarationWithIdentifier);
                             }
+                            if (declaration.tsPlusIndexTags && declaration.tsPlusIndexTags.length > 0) {
+                                file.tsPlusContext.index.push(declaration as VariableDeclarationWithIdentifier);
+                            }
+                            if (declaration.tsPlusPipeableIndexTags && declaration.tsPlusPipeableIndexTags.length > 0) {
+                                file.tsPlusContext.pipeableIndex.push(declaration as VariableDeclarationWithIdentifier);
+                            }
                         }
                     }
                     if (isFunctionDeclaration(statement) && statement.name) {
@@ -1279,6 +1285,9 @@ namespace ts {
                         if (statement.tsPlusIndexTags && statement.tsPlusIndexTags.length > 0) {
                             file.tsPlusContext.index.push(statement);
                         }
+                        if (statement.tsPlusPipeableIndexTags && statement.tsPlusPipeableIndexTags.length > 0) {
+                            file.tsPlusContext.pipeableIndex.push(statement);
+                        }
                     }
                 }
                 else {
@@ -7039,6 +7048,7 @@ namespace ts {
             const unifyTags: string[] = [];
             const macroTags: string[] = [];
             const indexTags: string[] = [];
+            const pipeableIndexTags: string[] = [];
             let isImplicit = false;
             for (const doc of jsDoc) {
                 if (doc.tags) {
@@ -7132,6 +7142,14 @@ namespace ts {
                                     indexTags.push(target);
                                     break;
                                 }
+                                case "pipeable-index": {
+                                    if (!target) {
+                                        parseErrorAt(tag.pos, tag.end - 1, Diagnostics.Annotation_of_an_index_extension_must_have_the_form_tsplus_index_typename);
+                                        break;
+                                    }
+                                    pipeableIndexTags.push(target);
+                                    break;
+                                }
                             }
                         }
                     }
@@ -7147,6 +7165,7 @@ namespace ts {
             (declaration as Mutable<FunctionDeclaration | VariableDeclaration>).tsPlusMacroTags = undefinedIfZeroLength(macroTags);
             (declaration as Mutable<FunctionDeclaration | VariableDeclaration>).tsPlusUnifyTags = undefinedIfZeroLength(unifyTags);
             (declaration as Mutable<FunctionDeclaration | VariableDeclaration>).tsPlusIndexTags = undefinedIfZeroLength(indexTags);
+            (declaration as Mutable<FunctionDeclaration | VariableDeclaration>).tsPlusPipeableIndexTags = undefinedIfZeroLength(pipeableIndexTags);
             if (isVariableDeclaration(declaration)) {
                 declaration.isTsPlusImplicit = isImplicit;
             }

src/compiler/transformers/tsplus.ts

@@ -250,6 +250,17 @@ namespace ts {
                 if (custom) {
                     const expression = visitNode(node.expression, visitor(source, traceInScope))
                     const argument = visitNode(node.argumentExpression, visitor(source, traceInScope))
+                    if (custom.signature && isTsPlusSignature(custom.signature) && custom.signature.tsPlusPipeable) {
+                        return context.factory.createCallExpression(
+                            context.factory.createCallExpression(
+                                getPathOfExtension(context, importer, custom, source, sourceFileUniqueNames),
+                                [],
+                                [argument]
+                            ),
+                            [],
+                            [expression]
+                        )
+                    }
                     return context.factory.createCallExpression(
                         getPathOfExtension(context, importer, custom, source, sourceFileUniqueNames),
                         [],

src/compiler/types.ts

@@ -1392,6 +1392,7 @@ namespace ts {
         tsPlusMacroTags?: string[];
         tsPlusUnifyTags?: string[];
         tsPlusIndexTags?: string[];
+        tsPlusPipeableIndexTags?: string[];
         tsPlusValidFluent?: boolean
     }
 
@@ -1610,6 +1611,7 @@ namespace ts {
         readonly tsPlusMacroTags?: string[];
         readonly tsPlusUnifyTags?: string[];
         readonly tsPlusIndexTags?: string[];
+        readonly tsPlusPipeableIndexTags?: string[];
         readonly tsPlusValidFluent?: boolean
         // TSPLUS END
     }
@@ -4194,7 +4196,8 @@ namespace ts {
             static: (VariableDeclarationWithIdentifier | FunctionDeclaration)[];
             getter: (VariableDeclarationWithIdentifier | FunctionDeclaration)[];
             unify: FunctionDeclaration[];
-            index: FunctionDeclaration[];
+            index: (VariableDeclarationWithIdentifier | FunctionDeclaration)[];
+            pipeableIndex: (VariableDeclarationWithIdentifier | FunctionDeclaration)[];
         }
         // TSPLUS EXTENSION END
     }
@@ -4970,7 +4973,7 @@ namespace ts {
         cloneSymbol(symbol: Symbol): Symbol
         getTextOfBinaryOp(kind: SyntaxKind): string | undefined
         getInstantiatedTsPlusSignature(declaration: Declaration, args: Expression[], checkMode: CheckMode | undefined): Signature
-        getIndexAccessExpressionCache(): ESMap<Node, { declaration: FunctionDeclaration, definition: SourceFile, exportName: string }>
+        getIndexAccessExpressionCache(): ESMap<Node, { signature: Signature, declaration: FunctionDeclaration | VariableDeclarationWithIdentifier, definition: SourceFile, exportName: string }>
         isTsPlusMacroCall<K extends string>(node: Node, macro: K): node is TsPlusMacroCallExpression<K>
         isTsPlusMacroGetter(node: Node, macro: string): boolean
         isClassCompanionReference(node: Expression): boolean