Implement special handling of subapplications when collecting MethodCallInfo (#9677)

Fixes on of the issues in #9354

Stale method call info for inner sub-application was causing additional argument placeholders on the node for certain expressions. Now it is fixed:
1. We only create function widget for the most top-level expression in the prefix application chain.
2. We reuse method call info from inner expressions, assuming it will be always correct for our purposes.

https://github.com/enso-org/enso/assets/6566674/91d2b4ba-a789-4c7b-b40c-f09ac45da7f0

Author: vitvakatu

app/gui2/src/components/GraphEditor/widgets/WidgetFunction.vue

@@ -17,12 +17,14 @@ import { useProjectStore, type NodeVisualizationConfiguration } from '@/stores/p
 import { entryQn } from '@/stores/suggestionDatabase/entry'
 import { assert, assertUnreachable } from '@/util/assert'
 import { Ast } from '@/util/ast'
+import type { AstId } from '@/util/ast/abstract'
 import {
   ArgumentApplication,
   ArgumentApplicationKey,
   ArgumentAst,
   ArgumentPlaceholder,
   getAccessOprSubject,
+  getMethodCallInfoRecursively,
   interpretCall,
 } from '@/util/callTree'
 import { partitionPoint } from '@/util/data/array'
@@ -37,12 +39,20 @@ const project = useProjectStore()
 
 provideFunctionInfo(
   proxyRefs({
-    callId: computed(() => props.input.value.id),
+    prefixCalls: computed(() => {
+      const ids = new Set<AstId>([props.input.value.id])
+      let ast: any = props.input.value
+      while (ast instanceof Ast.App) {
+        ids.add(ast.function.id)
+        ast = ast.function
+      }
+      return ids
+    }),
   }),
 )
 
 const methodCallInfo = computed(() => {
-  return graph.db.getMethodCallInfo(props.input.value.id)
+  return getMethodCallInfoRecursively(props.input.value, graph.db)
 })
 
 const interpreted = computed(() => {
@@ -306,18 +316,12 @@ export const widgetDefinition = defineWidget(
       // application with no arguments applied yet, but the application target is also an infix call.
       // In that case, the reentrant call method info must be ignored to not create an infinite loop,
       // and to resolve the infix call as its own application.
-      if (prevFunctionState?.callId === ast.id) return Score.Mismatch
+      // We only render the function widget on the application chain’s top-level.
+      if (prevFunctionState?.prefixCalls.has(ast.id)) return Score.Mismatch
 
       if (ast instanceof Ast.App || ast instanceof Ast.OprApp) return Score.Perfect
 
-      const info = db.getMethodCallInfo(ast.id)
-      if (
-        prevFunctionState != null &&
-        info?.partiallyApplied === true &&
-        ast instanceof Ast.Ident
-      ) {
-        return Score.Mismatch
-      }
+      const info = getMethodCallInfoRecursively(ast, db)
       return info != null ? Score.Perfect : Score.Mismatch
     },
   },

app/gui2/src/providers/functionInfo.ts

@@ -3,7 +3,8 @@ import type { AstId } from '@/util/ast/abstract.ts'
 import { identity } from '@vueuse/core'
 
 interface FunctionInfo {
-  callId: AstId | undefined
+  /** Ids of all nested prefix applications inside top-level expression (including the top-level). */
+  prefixCalls: Set<AstId>
 }
 
 export { injectFn as injectFunctionInfo, provideFn as provideFunctionInfo }

app/gui2/src/stores/graph/graphDatabase.ts

@@ -20,6 +20,11 @@ import type { ExternalId, SourceRange, VisualizationMetadata } from 'shared/yjsM
 import { isUuid, sourceRangeKey, visMetadataEquals } from 'shared/yjsModel'
 import { reactive, ref, type Ref } from 'vue'
 
+export interface MethodCallInfo {
+  methodCall: MethodCall
+  suggestion: SuggestionEntry
+}
+
 export interface BindingInfo {
   identifier: string
   usages: Set<AstId>
@@ -270,11 +275,7 @@ export class GraphDb {
     )
   }
 
-  getMethodCallInfo(
-    id: AstId,
-  ):
-    | { methodCall: MethodCall; suggestion: SuggestionEntry; partiallyApplied: boolean }
-    | undefined {
+  getMethodCallInfo(id: AstId): MethodCallInfo | undefined {
     const info = this.getExpressionInfo(id)
     if (info == null) return
     const payloadFuncSchema =
@@ -285,8 +286,7 @@ export class GraphDb {
     if (suggestionId == null) return
     const suggestion = this.suggestionDb.get(suggestionId)
     if (suggestion == null) return
-    const partiallyApplied = mathodCallEquals(methodCall, payloadFuncSchema)
-    return { methodCall, suggestion, partiallyApplied }
+    return { methodCall, suggestion }
   }
 
   getNodeColorStyle(id: NodeId): string {
@@ -524,7 +524,7 @@ const baseMockNode = {
   conditionalPorts: new Set(),
 } satisfies Partial<Node>
 
-function mathodCallEquals(a: MethodCall | undefined, b: MethodCall | undefined): boolean {
+export function mathodCallEquals(a: MethodCall | undefined, b: MethodCall | undefined): boolean {
   return (
     a === b ||
     (a != null &&

app/gui2/src/util/ast/__tests__/callTree.test.ts

@@ -1,14 +1,28 @@
 import * as widgetCfg from '@/providers/widgetRegistry/configuration'
-import { makeArgument, makeMethod, makeModuleMethod } from '@/stores/suggestionDatabase/entry'
+import { GraphDb } from '@/stores/graph/graphDatabase'
+import { ComputedValueRegistry, type ExpressionInfo } from '@/stores/project/computedValueRegistry'
+import { SuggestionDb } from '@/stores/suggestionDatabase'
+import {
+  makeArgument,
+  makeConstructor,
+  makeMethod,
+  makeModule,
+  makeModuleMethod,
+  makeType,
+  type SuggestionEntry,
+} from '@/stores/suggestionDatabase/entry'
 import { Ast } from '@/util/ast'
 import {
   ArgumentApplication,
   ArgumentAst,
   ArgumentPlaceholder,
+  getMethodCallInfoRecursively,
   interpretCall,
 } from '@/util/callTree'
 import { initializeFFI } from 'shared/ast/ffi'
+import type { ExpressionUpdatePayload, MethodCall } from 'shared/languageServerTypes'
 import { assert, expect, test } from 'vitest'
+import type { AstId } from '../abstract'
 
 await initializeFFI()
 
@@ -107,6 +121,150 @@ test.each`
   },
 )
 
+interface TestCase {
+  description: string
+  code: string
+  /** Index of sub-application for which the MethodCallInfo is available. */
+  subapplicationIndex: number
+  /** Not applied arguments in available MethodCallInfo. */
+  notAppliedArguments: number[]
+  /** Not applied arguments expected for the whole expression. */
+  expectedNotAppliedArguments: number[]
+}
+
+test.each([
+  {
+    description: 'Base case',
+    code: 'Aggregate_Column.Sum',
+    subapplicationIndex: 0,
+    notAppliedArguments: [0, 1],
+    expectedNotAppliedArguments: [0, 1],
+  },
+  {
+    description: '1 arg, info on most inner subapplication.',
+    code: 'Aggregate_Column.Sum x',
+    subapplicationIndex: 1,
+    notAppliedArguments: [0, 1],
+    expectedNotAppliedArguments: [1],
+  },
+  {
+    description: '2 args, info on most inner subapplication.',
+    code: 'Aggregate_Column.Sum x y',
+    subapplicationIndex: 2,
+    notAppliedArguments: [0, 1],
+    expectedNotAppliedArguments: [],
+  },
+  {
+    description: '2 args, info on inner subapplication.',
+    code: 'Aggregate_Column.Sum x y',
+    subapplicationIndex: 1,
+    notAppliedArguments: [1],
+    expectedNotAppliedArguments: [],
+  },
+  {
+    description: '2 args, notAppliedArguments are incorrectly empty.',
+    code: 'Aggregate_Column.Sum x y',
+    subapplicationIndex: 2,
+    notAppliedArguments: [],
+    expectedNotAppliedArguments: [],
+  },
+  {
+    description: '1 arg, notAppliedArguments unsorted.',
+    code: 'Aggregate_Column.Sum x',
+    subapplicationIndex: 1,
+    notAppliedArguments: [1, 0],
+    expectedNotAppliedArguments: [1],
+  },
+  {
+    description: '1 named arg.',
+    code: 'Aggregate_Column.Sum as=x',
+    subapplicationIndex: 1,
+    notAppliedArguments: [0, 1],
+    expectedNotAppliedArguments: [0],
+  },
+  {
+    description: '2 named args.',
+    code: 'Aggregate_Column.Sum as=x column=y',
+    subapplicationIndex: 2,
+    notAppliedArguments: [0, 1],
+    expectedNotAppliedArguments: [],
+  },
+  {
+    description: '1 wrongly named arg.',
+    code: 'Aggregate_Column.Sum bla=x',
+    subapplicationIndex: 1,
+    notAppliedArguments: [0, 1],
+    expectedNotAppliedArguments: [0, 1],
+  },
+  {
+    description: '1 named & 1 unnamed args.',
+    code: 'Aggregate_Column.Sum as=x y',
+    subapplicationIndex: 2,
+    notAppliedArguments: [0, 1],
+    expectedNotAppliedArguments: [],
+  },
+] as TestCase[])(
+  'Getting MethodCallInfo for sub-applications: $description',
+  ({ code, subapplicationIndex, notAppliedArguments, expectedNotAppliedArguments }: TestCase) => {
+    const { db, expectedMethodCall, expectedSuggestion, setExpressionInfo } =
+      prepareMocksForGetMethodCallTest()
+    const ast = Ast.parse(code)
+    const subApplication = nthSubapplication(ast, subapplicationIndex)
+    assert(subApplication)
+    db.updateExternalIds(ast)
+    setExpressionInfo(subApplication.id, {
+      typename: undefined,
+      methodCall: { ...expectedMethodCall, notAppliedArguments },
+      payload: { type: 'Pending' } as ExpressionUpdatePayload,
+      profilingInfo: [],
+    })
+
+    const info = getMethodCallInfoRecursively(ast, db)
+    expect(info?.methodCall).toEqual({
+      ...expectedMethodCall,
+      notAppliedArguments: expectedNotAppliedArguments,
+    })
+    expect(info?.suggestion).toEqual(expectedSuggestion)
+  },
+)
+
+function prepareMocksForGetMethodCallTest(): {
+  db: GraphDb
+  expectedMethodCall: MethodCall
+  expectedSuggestion: SuggestionEntry
+  setExpressionInfo: (id: AstId, info: ExpressionInfo) => void
+} {
+  const suggestionDb = new SuggestionDb()
+  suggestionDb.set(1, makeModule('Standard.Table.Aggregate_Column'))
+  suggestionDb.set(2, makeType('Standard.Table.Aggregate_Column.Aggregate_Column'))
+  const con = makeConstructor('Standard.Table.Aggregate_Column.Aggregate_Column.Sum')
+  con.arguments = [makeArgument('column', 'Any'), makeArgument('as', 'Any')]
+  suggestionDb.set(3, con)
+  const expectedSuggestion = suggestionDb.get(3)!
+  const registry = ComputedValueRegistry.Mock()
+  const db = GraphDb.Mock(registry, suggestionDb)
+  const expectedMethodCall = {
+    methodPointer: {
+      module: 'Standard.Table.Aggregate_Column',
+      definedOnType: 'Standard.Table.Aggregate_Column.Aggregate_Column',
+      name: 'Sum',
+    },
+    notAppliedArguments: [0, 1],
+  }
+  const setExpressionInfo = (id: AstId, info: ExpressionInfo) =>
+    registry.db.set(db.idToExternal(id)!, info)
+  return { db, expectedMethodCall, expectedSuggestion, setExpressionInfo }
+}
+
+/** Nth sub-application of the Ast.App call chain. 0th is a `root` itself. */
+function nthSubapplication(root: Ast.Ast, n: number): Ast.Ast | undefined {
+  let current: Ast.Ast | undefined = root
+  for (let i = 0; i < n; i++) {
+    current = current instanceof Ast.App ? current.function : undefined
+  }
+  return current
+}
+
 function printArgPattern(application: ArgumentApplication | Ast.Ast) {
   const parts: string[] = []
   let current: ArgumentApplication['target'] = application

app/gui2/src/util/callTree.ts

@@ -3,6 +3,7 @@ import { WidgetInput } from '@/providers/widgetRegistry'
 import type { WidgetConfiguration } from '@/providers/widgetRegistry/configuration'
 import * as widgetCfg from '@/providers/widgetRegistry/configuration'
 import { DisplayMode } from '@/providers/widgetRegistry/configuration'
+import type { GraphDb, MethodCallInfo } from '@/stores/graph/graphDatabase'
 import type { SuggestionEntry, SuggestionEntryArgument } from '@/stores/suggestionDatabase/entry'
 import { Ast } from '@/util/ast'
 import { findLastIndex, tryGetIndex } from '@/util/data/array'
@@ -409,6 +410,54 @@ export function getAccessOprSubject(app: Ast.Ast): Ast.Ast | undefined {
   if (app instanceof Ast.PropertyAccess) return app.lhs
 }
 
+/** Same as {@link GraphDb.getMethodCallInfo} but with a special handling for nested Applications.
+ * Sometimes we receive MethodCallInfo for inner sub-applications of the expression,
+ * and we want to reuse it for outer application expressions when adding new arguments to the call.
+ * It requires adjusting `notAppliedArguments` array, but otherwise is a straightforward recursive call.
+ * We stop recursion at any not-application AST. We expect that a subexpression’s info is only valid if it is a part of the prefix application chain.
+ * We also don’t consider infix applications here, as using them inside a prefix chain would require additional syntax (like parenthesis). */
+export function getMethodCallInfoRecursively(
+  ast: Ast.Ast,
+  db: GraphDb,
+): MethodCallInfo | undefined {
+  let appliedArgs = 0
+  const appliedNamedArgs: string[] = []
+  for (;;) {
+    const info = db.getMethodCallInfo(ast.id)
+    if (info) {
+      // There is an info available! Stop the recursion and adjust `notAppliedArguments`.
+      // Indices of all named arguments applied so far.
+      const appliedNamed =
+        appliedNamedArgs.length > 0 ?
+          info.suggestion.arguments
+            .map((arg, index) => (appliedNamedArgs.includes(arg.name) ? index : -1))
+            .filter((i) => i !== -1)
+        : []
+      const withoutNamed = info.methodCall.notAppliedArguments.filter(
+        (idx) => !appliedNamed.includes(idx),
+      )
+      return {
+        methodCall: {
+          ...info.methodCall,
+          notAppliedArguments: withoutNamed.sort().slice(appliedArgs),
+        },
+        suggestion: info.suggestion,
+      }
+    }
+    // No info, continue recursion to the next sub-application AST.
+    if (ast instanceof Ast.App) {
+      if (ast.argumentName) {
+        appliedNamedArgs.push(ast.argumentName.code())
+      } else {
+        appliedArgs += 1
+      }
+      ast = ast.function
+    } else {
+      break
+    }
+  }
+}
+
 export const ArgumentApplicationKey: unique symbol = Symbol('ArgumentApplicationKey')
 export const ArgumentInfoKey: unique symbol = Symbol('ArgumentInfoKey')
 declare module '@/providers/widgetRegistry' {