improve completion and onhover in VSCode for decorator and extended model

.chronus/changes/add-completion-for-decorator-model-arg-2024-4-6-16-27-17.md

@@ -0,0 +1,7 @@
+---
+changeKind: feature
+packages:
+  - "@typespec/compiler"
+---
+
+Support completion for Model Expression as decorator argument value and Extended Model Properties

packages/compiler/src/core/checker.ts

@@ -16,6 +16,7 @@ import { isStringTemplateSerializable } from "./helpers/string-template-utils.js
 import { createDiagnostic } from "./messages.js";
 import {
   exprIsBareIdentifier,
+  getFirstAncestor,
   getIdentifierContext,
   hasParseError,
   visitChildren,
@@ -2055,6 +2056,15 @@ export function createChecker(program: Program): Checker {
     const { node, kind } = getIdentifierContext(id);
 
     switch (kind) {
+      case IdentifierKind.ModelExpressionProperty:
+        const model = getReferencedModelFromDecoratorArgument(node as ModelPropertyNode);
+        if (model) {
+          sym = getMemberSymbol(model.node!.symbol, id.sv);
+        } else {
+          return undefined;
+        }
+        break;
+      case IdentifierKind.ModelStatementProperty:
       case IdentifierKind.Declaration:
         if (node.symbol && (!isTemplatedNode(node) || mapper === undefined)) {
           sym = getMergedSymbol(node.symbol);
@@ -2133,6 +2143,63 @@ export function createChecker(program: Program): Checker {
     return (resolved?.declarations.filter((n) => isTemplatedNode(n)) ?? []) as TemplateableNode[];
   }
 
+  function getReferencedModelFromDecoratorArgument(
+    modelExpressionProperty: ModelPropertyNode
+  ): Model | undefined {
+    if (modelExpressionProperty.parent?.kind !== SyntaxKind.ModelExpression) {
+      return undefined;
+    }
+    const path: string[] = [];
+    const decArgNode = getFirstAncestor(modelExpressionProperty, (n) => {
+      if (
+        n.kind === SyntaxKind.ModelExpression &&
+        n.parent?.kind === SyntaxKind.DecoratorExpression
+      ) {
+        return true;
+      } else if (n.kind === SyntaxKind.ModelProperty) {
+        path.unshift(n.id.sv);
+      }
+      return false;
+    });
+    const decNode = decArgNode?.parent;
+    if (!decArgNode || decNode?.kind !== SyntaxKind.DecoratorExpression) {
+      return undefined;
+    }
+
+    const decSym = program.checker.resolveIdentifier(
+      decNode.target.kind === SyntaxKind.MemberExpression ? decNode.target.id : decNode.target
+    );
+    if (!decSym) {
+      return undefined;
+    }
+
+    const decDecl: DecoratorDeclarationStatementNode | undefined = decSym.declarations.find(
+      (x): x is DecoratorDeclarationStatementNode =>
+        x.kind === SyntaxKind.DecoratorDeclarationStatement
+    );
+    if (!decDecl) {
+      return undefined;
+    }
+
+    const decType = program.checker.getTypeForNode(decDecl);
+    compilerAssert(decType.kind === "Decorator", "Expected type to be a decorator.");
+
+    const argIndex = decNode.arguments.findIndex((n) => n === decArgNode);
+    if (argIndex >= decType.parameters.length) {
+      return undefined;
+    }
+    const decArg = decType.parameters[argIndex];
+
+    let parentType: Model | undefined = decArg.type as Model;
+    if (parentType?.kind !== "Model") return undefined;
+    for (const seg of path) {
+      parentType = parentType?.properties.get(seg)?.type as Model;
+      if (parentType?.kind !== "Model") return undefined;
+    }
+
+    return parentType;
+  }
+
   function resolveCompletions(identifier: IdentifierNode): Map<string, TypeSpecCompletionItem> {
     const completions = new Map<string, TypeSpecCompletionItem>();
     const { kind, node: ancestor } = getIdentifierContext(identifier);
@@ -2142,6 +2209,8 @@ export function createChecker(program: Program): Checker {
       case IdentifierKind.Decorator:
       case IdentifierKind.Function:
       case IdentifierKind.TypeReference:
+      case IdentifierKind.ModelExpressionProperty:
+      case IdentifierKind.ModelStatementProperty:
         break; // supported
       case IdentifierKind.Other:
         return completions; // not implemented
@@ -2166,7 +2235,40 @@ export function createChecker(program: Program): Checker {
         compilerAssert(false, "Unreachable");
     }
 
-    if (identifier.parent && identifier.parent.kind === SyntaxKind.MemberExpression) {
+    if (kind === IdentifierKind.ModelStatementProperty) {
+      const model = ancestor.parent as ModelStatementNode;
+      const modelType = program.checker.getTypeForNode(model) as Model;
+      const baseType = modelType.baseModel;
+      const baseNode = baseType?.node;
+      if (!baseNode) {
+        return completions;
+      }
+      for (const prop of baseType.properties.values()) {
+        if (identifier.sv === prop.name || !modelType.properties.has(prop.name)) {
+          const sym = getMemberSymbol(baseNode.symbol, prop.name);
+          if (sym) {
+            addCompletion(prop.name, sym);
+          }
+        }
+      }
+    } else if (kind === IdentifierKind.ModelExpressionProperty) {
+      // When a ModelExpression is used as decorator model argument, we should be
+      // able to figure out completions from the decorator argument's definition
+      const model = getReferencedModelFromDecoratorArgument(ancestor as ModelPropertyNode);
+      if (!model) return completions;
+      const editingPropertyNames = (ancestor.parent as ModelExpressionNode)?.properties
+        .filter((p) => p.kind === SyntaxKind.ModelProperty && p.id.sv !== identifier.sv)
+        .map((p) => (p as ModelPropertyNode).id.sv);
+
+      for (const prop of walkPropertiesInherited(model)) {
+        if (!editingPropertyNames?.includes(prop.name)) {
+          const sym = getMemberSymbol(model.node!.symbol, prop.name);
+          if (sym) {
+            addCompletion(prop.name, sym);
+          }
+        }
+      }
+    } else if (identifier.parent && identifier.parent.kind === SyntaxKind.MemberExpression) {
       let base = resolveTypeReferenceSym(identifier.parent.base, undefined, false);
       if (base) {
         if (base.flags & SymbolFlags.Alias) {
@@ -2271,6 +2373,9 @@ export function createChecker(program: Program): Checker {
 
     function shouldAddCompletion(sym: Sym): boolean {
       switch (kind) {
+        case IdentifierKind.ModelExpressionProperty:
+        case IdentifierKind.ModelStatementProperty:
+          return !!(sym.flags & SymbolFlags.ModelProperty);
         case IdentifierKind.Decorator:
           // Only return decorators and namespaces when completing decorator
           return !!(sym.flags & (SymbolFlags.Decorator | SymbolFlags.Namespace));

packages/compiler/src/core/parser.ts

@@ -3560,6 +3560,23 @@ export function getIdentifierContext(id: IdentifierNode): IdentifierContext {
     case SyntaxKind.TemplateArgument:
       kind = IdentifierKind.TemplateArgument;
       break;
+    case SyntaxKind.ModelProperty:
+      switch (node.parent?.kind) {
+        case SyntaxKind.ModelExpression:
+          kind = IdentifierKind.ModelExpressionProperty;
+          break;
+        case SyntaxKind.ModelStatement:
+          kind = IdentifierKind.ModelStatementProperty;
+          break;
+        default:
+          compilerAssert("false", "ModelProperty with unexpected parent kind.");
+          kind =
+            (id.parent as DeclarationNode).id === id
+              ? IdentifierKind.Declaration
+              : IdentifierKind.Other;
+          break;
+      }
+      break;
     default:
       kind =
         (id.parent as DeclarationNode).id === id

packages/compiler/src/core/types.ts

@@ -1649,6 +1649,8 @@ export enum IdentifierKind {
   Function,
   Using,
   Declaration,
+  ModelExpressionProperty,
+  ModelStatementProperty,
   Other,
 }
 

packages/compiler/src/server/completion.ts

@@ -7,18 +7,25 @@ import {
   MarkupKind,
   TextEdit,
 } from "vscode-languageserver";
+import { TextDocument } from "vscode-languageserver-textdocument";
+import { CharCode } from "../core/charcode.js";
 import { getDeprecationDetails } from "../core/deprecation.js";
 import {
   CompilerHost,
+  DecoratorExpressionNode,
   IdentifierNode,
+  ModelExpressionNode,
+  ModelStatementNode,
   Node,
+  NodeFlags,
   NodePackage,
   Program,
   StringLiteralNode,
   SymbolFlags,
   SyntaxKind,
   Type,
   TypeSpecScriptNode,
+  getFirstAncestor,
 } from "../core/index.js";
 import {
   getAnyExtensionFromPath,
@@ -34,6 +41,7 @@ import { getSymbolDetails } from "./type-details.js";
 export type CompletionContext = {
   program: Program;
   params: CompletionParams;
+  document: TextDocument;
   file: TypeSpecScriptNode;
   completions: CompletionList;
 };
@@ -64,6 +72,31 @@ export async function resolveCompletion(
           await addImportCompletion(context, node);
         }
         break;
+      case SyntaxKind.ModelStatement:
+        // we can provide completions for overrides when the model has base model
+        if (node.extends) {
+          addModelCompletion(context, node);
+        }
+        break;
+      case SyntaxKind.ModelExpression:
+        // check for following scenario:
+        // a model expresssion as a decorator argument like `@dec({ | })`
+        let decNode: DecoratorExpressionNode | undefined = undefined;
+        if (node.parent?.kind === SyntaxKind.DecoratorExpression) {
+          decNode = node.parent;
+        } else {
+          const argNode = getFirstAncestor(
+            node,
+            (n) =>
+              n.kind === SyntaxKind.ModelExpression &&
+              n.parent?.kind === SyntaxKind.DecoratorExpression
+          );
+          decNode = argNode?.parent as DecoratorExpressionNode;
+        }
+        if (decNode) {
+          addModelCompletion(context, node);
+        }
+        break;
     }
   }
 
@@ -229,6 +262,33 @@ async function addRelativePathCompletion(
   }
 }
 
+function addModelCompletion(
+  context: CompletionContext,
+  node: ModelStatementNode | ModelExpressionNode
+) {
+  // skip the scenario like `{ ... }|`
+  const cur = context.document.offsetAt(context.params.position);
+  if (cur === node.end) {
+    const endChar = context.file.file.text.charCodeAt(cur - 1);
+    if (endChar === CharCode.CloseBrace) {
+      return;
+    }
+  }
+  // create a fake identifier node in the model to further resolve the completions
+  const fakeProp = {
+    kind: SyntaxKind.ModelProperty,
+    flags: NodeFlags.None,
+    parent: node,
+  };
+  const fakeId = {
+    kind: SyntaxKind.Identifier,
+    sv: "",
+    flags: NodeFlags.None,
+    parent: fakeProp,
+  };
+  addIdentifierCompletion(context, fakeId as IdentifierNode);
+}
+
 /**
  * Add completion options for an identifier.
  */

packages/compiler/src/server/serverlib.ts

@@ -676,6 +676,7 @@ export function createServer(host: ServerHost): Server {
         {
           program,
           file: script,
+          document,
           completions,
           params,
         },