Repository: elm-tooling/elm-language-server (ts)

Issue:
Offer to create a record field if it is missing
This would be especially useful with the `Model` type, where it can be common to write out a record field before it exists. We should be able to infer the expected type of that field and create it. 

Pull request description:
Add code action to create missing record field
Closes #477. 

Patch:
diff --git a/src/providers/codeAction/addMissingRecordFieldCodeAction.ts b/src/providers/codeAction/addMissingRecordFieldCodeAction.ts
new file mode 100644
index 0000000..75d3ac8
--- /dev/null
+++ b/src/providers/codeAction/addMissingRecordFieldCodeAction.ts
@@ -0,0 +1,150 @@
+import { CodeAction, Range, TextEdit } from "vscode-languageserver";
+import { SyntaxNode } from "web-tree-sitter";
+import { CodeActionProvider } from "..";
+import { ITreeContainer } from "../../forest";
+import { getSpaces } from "../../util/refactorEditUtils";
+import { TreeUtils } from "../../util/treeUtils";
+import { Diagnostics } from "../../util/types/diagnostics";
+import { TypeChecker } from "../../util/types/typeChecker";
+import { Type } from "../../util/types/typeInference";
+import { ICodeActionParams } from "../paramsExtensions";
+
+const errorCodes = [Diagnostics.RecordField.code];
+const fixId = "add_missing_record_field";
+
+CodeActionProvider.registerCodeAction({
+  errorCodes,
+  fixId,
+  getCodeActions: (params: ICodeActionParams) => {
+    const edits = getEdits(params, params.range);
+
+    if (Object.values(edits).some((edit) => edit.length > 0)) {
+      const single = Object.values(edits).every((edit) => edit.length === 1);
+
+      return [
+        CodeActionProvider.getCodeAction(
+          params,
+          `Create missing record ${single ? "field" : "fields"}`,
+          edits,
+        ),
+      ];
+    }
+
+    return [];
+  },
+  getFixAllCodeAction: (params: ICodeActionParams): CodeAction | undefined => {
+    return CodeActionProvider.getFixAllCodeAction(
+      "Create all missing record fields",
+      params,
+      errorCodes,
+      fixId,
+      () => {
+        // Not used
+      },
+      (changes, diagnostic) => {
+        mergeChanges(changes, getEdits(params, diagnostic.range));
+      },
+    );
+  },
+});
+
+function getEdits(
+  params: ICodeActionParams,
+  range: Range,
+): { [uri: string]: TextEdit[] } {
+  const nodeAtPosition = TreeUtils.getNamedDescendantForRange(
+    params.sourceFile,
+    range,
+  );
+
+  const nodeParent = nodeAtPosition.parent;
+  if (
+    nodeAtPosition.type === "lower_case_identifier" &&
+    (nodeParent?.type === "field_access_expr" || nodeParent?.type === "field")
+  ) {
+    const checker = params.program.getTypeChecker();
+    const fieldName = nodeAtPosition.text;
+
+    let target =
+      nodeParent.type === "field_access_expr"
+        ? nodeParent?.childForFieldName("target")
+        : nodeParent.type === "field"
+        ? nodeParent.parent
+        : nodeAtPosition;
+
+    // Adjust for parenthesis expr. Will need to change when we handle it better in inference
+    if (target?.type === "parenthesized_expr") {
+      target = target.namedChildren[1];
+    }
+
+    if (target) {
+      const expectedType = checker.findType(
+        nodeParent.type === "field_access_expr" ? nodeParent : nodeAtPosition,
+      );
+
+      return createFields(
+        [[fieldName, expectedType]],
+        target,
+        checker,
+        params.sourceFile,
+      );
+    }
+  }
+
+  return {};
+}
+
+function createFields(
+  fields: [string, Type][],
+  targetRecord: SyntaxNode,
+  checker: TypeChecker,
+  sourceFile: ITreeContainer,
+): { [uri: string]: TextEdit[] } {
+  const type = checker.findType(targetRecord);
+
+  if (type.nodeType === "Record") {
+    const fieldNames = Object.keys(type.fields);
+    const lastFieldRef = TreeUtils.findFieldReference(
+      type,
+      fieldNames[fieldNames.length - 1],
+    );
+
+    if (lastFieldRef) {
+      const useNewLine =
+        lastFieldRef.node.parent?.startPosition.row !==
+        lastFieldRef.node.parent?.endPosition.row;
+      const indendation = getSpaces(lastFieldRef.node.startPosition.column - 2);
+
+      const edits = fields.map(([fieldName, fieldType]) => {
+        const typeString: string = checker.typeToString(fieldType, sourceFile);
+
+        return TextEdit.insert(
+          {
+            line: lastFieldRef?.node.endPosition.row,
+            character: lastFieldRef?.node.endPosition.column,
+          },
+          `${
+            useNewLine ? "\n" + indendation : ""
+          }, ${fieldName} : ${typeString}`,
+        );
+      });
+
+      return { [lastFieldRef.uri]: edits };
+    }
+  }
+
+  return {};
+}
+
+function mergeChanges(
+  a: { [uri: string]: TextEdit[] },
+  b: { [uri: string]: TextEdit[] },
+): void {
+  Object.entries(b).forEach(([uri, edits]) => {
+    if (a[uri]) {
+      a[uri].push(...edits);
+    } else {
+      a[uri] = edits;
+    }
+  });
+}
diff --git a/src/providers/codeAction/index.ts b/src/providers/codeAction/index.ts
index 6270240..22fa282 100644
--- a/src/providers/codeAction/index.ts
+++ b/src/providers/codeAction/index.ts
@@ -5,3 +5,4 @@ import "./extractFunctionCodeAction";
 import "./exposeUnexposeCodeAction";
 import "./moveFunctionCodeAction";
 import "./installPackageCodeAction";
+import "./addMissingRecordFieldCodeAction";
diff --git a/src/providers/codeActionProvider.ts b/src/providers/codeActionProvider.ts
index ccf3542..134757c 100644
--- a/src/providers/codeActionProvider.ts
+++ b/src/providers/codeActionProvider.ts
@@ -161,10 +161,12 @@ export class CodeActionProvider {
   public static getCodeAction(
     params: ICodeActionParams,
     title: string,
-    edits: TextEdit[],
+    edits: TextEdit[] | { [uri: string]: TextEdit[] },
     command?: Command,
   ): CodeAction {
-    const changes = { [params.sourceFile.uri]: edits };
+    const changes = Array.isArray(edits)
+      ? { [params.sourceFile.uri]: edits }
+      : edits;
     return {
       title,
       kind: CodeActionKind.QuickFix,
@@ -180,16 +182,26 @@ export class CodeActionProvider {
     errorCodes: string[],
     fixId: string,
     callback: (edits: TextEdit[], diagnostic: LspDiagnostic) => void,
+    callbackChanges?: (
+      edits: { [uri: string]: TextEdit[] },
+      diagnostic: LspDiagnostic,
+    ) => void,
   ): CodeAction {
     const edits: TextEdit[] = [];
-    const changes = {
-      [params.sourceFile.uri]: edits,
-    };
+    const changes = callbackChanges
+      ? {}
+      : {
+          [params.sourceFile.uri]: edits,
+        };
 
     const diagnostics: LspDiagnostic[] = [];
     CodeActionProvider.forEachDiagnostic(params, errorCodes, (diagnostic) => {
       diagnostics.push(diagnostic);
-      callback(edits, diagnostic);
+      if (callbackChanges) {
+        callbackChanges(changes, diagnostic);
+      } else {
+        callback(edits, diagnostic);
+      }
     });
 
     return {
diff --git a/src/providers/completionProvider.ts b/src/providers/completionProvider.ts
index 6c3287e..ae5cef0 100644
--- a/src/providers/completionProvider.ts
+++ b/src/providers/completionProvider.ts
@@ -141,30 +141,6 @@ export class CompletionProvider {
       ) {
         // Don't complete in comments
         return [];
-      } else if (
-        nodeAtPosition.parent?.type === "lower_pattern" ||
-        nodeAtPosition.parent?.type === "record_pattern"
-      ) {
-        if (
-          nodeAtPosition.parent.parent?.type === "record_pattern" ||
-          nodeAtPosition.parent.type === "record_pattern"
-        ) {
-          return this.getRecordCompletionsUsingInference(
-            checker,
-            TreeUtils.findParentOfType("record_pattern", nodeAtPosition) ??
-              nodeAtPosition.parent,
-            replaceRange,
-          );
-        } else {
-          // Don't complete on pattern names
-          return [];
-        }
-      } else if (
-        nodeAtPosition.parent?.type === "as_clause" &&
-        nodeAtPosition.parent.parent?.type === "import_clause"
-      ) {
-        // Don't complete on import alias
-        return [];
       } else if (
         isAtStartOfLine &&
         nodeAtLineBefore.type === "lower_case_identifier" &&
@@ -394,6 +370,7 @@ export class CompletionProvider {
         ...this.getCompletionsFromOtherFile(
           checker,
           treeContainer,
+          params.textDocument.uri,
           replaceRange,
           targetWord,
         ),
@@ -549,6 +526,7 @@ export class CompletionProvider {
   private getCompletionsFromOtherFile(
     checker: TypeChecker,
     treeContainer: ITreeContainer,
+    uri: string,
     range: Range,
     inputText: string,
   ): CompletionItem[] {
@@ -573,9 +551,9 @@ export class CompletionProvider {
       const dotIndex = label.lastIndexOf(".");
       const valuePart = label.slice(dotIndex + 1);
 
-      const importNode = checker.findImportModuleNameNode(
+      const importNode = TreeUtils.findImportClauseByName(
+        treeContainer.tree,
         element.fromModuleName,
-        treeContainer,
       );
 
       // Check if a value is already imported for this module using the exposing list
@@ -587,7 +565,6 @@ export class CompletionProvider {
       // Try to determine if just the value is being typed
       if (
         !valuesAlreadyExposed &&
-        element.fromModuleName !== "Basics" &&
         valuePart.toLowerCase().startsWith(inputText.toLowerCase())
       ) {
         filterText = valuePart;
@@ -896,13 +873,6 @@ export class CompletionProvider {
     return this.createCompletion(options);
   }
 
-  private createVariableCompletion(
-    options: ICompletionOptions,
-  ): CompletionItem {
-    options.kind = CompletionItemKind.Variable;
-    return this.createCompletion(options);
-  }
-
   private createFieldOrParameterCompletion(
     markdownDocumentation: string | undefined,
     label: string,
@@ -996,10 +966,13 @@ export class CompletionProvider {
       if (node.parent.type === "let_in_expr") {
         node.parent.children.forEach((nodeToProcess) => {
           if (
+            nodeToProcess &&
             nodeToProcess.type === "value_declaration" &&
-            nodeToProcess.firstNamedChild?.type ===
+            nodeToProcess.firstNamedChild !== null &&
+            nodeToProcess.firstNamedChild.type ===
               "function_declaration_left" &&
-            nodeToProcess.firstNamedChild.firstNamedChild?.type ===
+            nodeToProcess.firstNamedChild.firstNamedChild !== null &&
+            nodeToProcess.firstNamedChild.firstNamedChild.type ===
               "lower_case_identifier"
           ) {
             const markdownDocumentation = HintHelper.createHintFromDefinitionInLet(
@@ -1016,26 +989,30 @@ export class CompletionProvider {
           }
         });
       }
-      if (node.parent.type === "case_of_branch") {
-        const pattern = node.parent.childForFieldName("pattern");
-
-        if (pattern) {
-          const caseBranchVariableNodes = pattern.descendantsOfType(
-            "lower_pattern",
-          );
-          if (caseBranchVariableNodes) {
-            caseBranchVariableNodes.forEach((a) => {
-              const markdownDocumentation = HintHelper.createHintFromDefinitionInCaseBranch();
-              result.push(
-                this.createVariableCompletion({
-                  markdownDocumentation,
-                  label: a.text,
-                  range,
-                  sortPrefix,
-                }),
-              );
-            });
-          }
+      if (
+        node.parent.type === "case_of_branch" &&
+        node.parent.firstNamedChild &&
+        node.parent.firstNamedChild.type === "pattern" &&
+        node.parent.firstNamedChild.firstNamedChild &&
+        node.parent.firstNamedChild.firstNamedChild.type === "union_pattern" &&
+        node.parent.firstNamedChild.firstNamedChild.childCount > 1 // Ignore cases of case branches with no params
+      ) {
+        const caseBranchVariableNodes = TreeUtils.findAllNamedChildrenOfType(
+          "lower_pattern",
+          node.parent.firstNamedChild.firstNamedChild,
+        );
+        if (caseBranchVariableNodes) {
+          caseBranchVariableNodes.forEach((a) => {
+            const markdownDocumentation = HintHelper.createHintFromDefinitionInCaseBranch();
+            result.push(
+              this.createFunctionCompletion({
+                markdownDocumentation,
+                label: a.text,
+                range,
+                sortPrefix,
+              }),
+            );
+          });
         }
       }
       if (
@@ -1049,12 +1026,11 @@ export class CompletionProvider {
               child,
             );
             result.push(
-              this.createVariableCompletion({
+              this.createFieldOrParameterCompletion(
                 markdownDocumentation,
-                label: child.text,
+                child.text,
                 range,
-                sortPrefix,
-              }),
+              ),
             );
 
             const annotationTypeNode = TreeUtils.getTypeOrTypeAliasOfFunctionParameter(
@@ -1106,8 +1082,14 @@ export class CompletionProvider {
     const sourceFile = forest.getByUri(uri);
 
     if (sourceFile) {
-      const checker = program.getTypeChecker();
-      const allImportedValues = checker.getAllImports(sourceFile);
+      const allImportedValues = program
+        .getTypeChecker()
+        .getAllImports(sourceFile);
+
+      const importedModules =
+        TreeUtils.findAllImportClauseNodes(sourceFile.tree)?.map(
+          (n) => TreeUtils.findFirstNamedChildOfType("upper_case_qid", n)?.text,
+        ) ?? [];
 
       const cached = possibleImportsCache.get(uri);
       const possibleImports =
@@ -1123,13 +1105,13 @@ export class CompletionProvider {
         .filter(
           (possibleImport): boolean =>
             !allImportedValues.get(
-              possibleImport.value,
+              possibleImport.valueToImport ?? possibleImport.value,
               (imp) => imp.fromModuleName === possibleImport.module,
             ),
         )
         .sort((a, b) => {
-          const aValue = a.value.toLowerCase();
-          const bValue = b.value.toLowerCase();
+          const aValue = (a.valueToImport ?? a.value).toLowerCase();
+          const bValue = (b.valueToImport ?? b.value).toLowerCase();
 
           filterText = filterText.toLowerCase();
 
@@ -1150,14 +1132,8 @@ export class CompletionProvider {
             } else if (!aMatches && bMatches) {
               return 1;
             } else {
-              const aModuleImported = !!checker.findImportModuleNameNode(
-                a.module,
-                sourceFile,
-              );
-              const bModuleImported = !!checker.findImportModuleNameNode(
-                b.module,
-                sourceFile,
-              );
+              const aModuleImported = importedModules.includes(a.module);
+              const bModuleImported = importedModules.includes(b.module);
 
               if (aModuleImported && !bModuleImported) {
                 return -1;
diff --git a/src/providers/diagnostics/elmLsDiagnostics.ts b/src/providers/diagnostics/elmLsDiagnostics.ts
index 9100f6b..c85cb1a 100644
--- a/src/providers/diagnostics/elmLsDiagnostics.ts
+++ b/src/providers/diagnostics/elmLsDiagnostics.ts
@@ -13,14 +13,7 @@ import {
   TextEdit,
 } from "vscode-languageserver";
 import { URI } from "vscode-uri";
-import {
-  Language,
-  Parser,
-  Query,
-  QueryResult,
-  SyntaxNode,
-  Tree,
-} from "web-tree-sitter";
+import { Language, Parser, Query, SyntaxNode, Tree } from "web-tree-sitter";
 import { IElmWorkspace } from "../../elmWorkspace";
 import { PositionUtil } from "../../positionUtil";
 import { ElmWorkspaceMatcher } from "../../util/elmWorkspaceMatcher";
@@ -29,7 +22,6 @@ import { TreeUtils } from "../../util/treeUtils";
 import { Utils } from "../../util/utils";
 import { IDiagnostic } from "./diagnosticsProvider";
 import * as path from "path";
-import { SyntaxNodeMap } from "../../util/types/syntaxNodeMap";
 
 interface IElmAnalyseJson {
   checks?: {
@@ -72,7 +64,6 @@ export class ElmLsDiagnostics {
   private elmAnalyseJson = new Map<string, IElmAnalyseJson>();
 
   private readonly exposedValuesAndTypesQuery: Query;
-  private readonly exposedValueAndTypeUsagesQuery: Query;
   private readonly moduleImportsQuery: Query;
   private readonly moduleReferencesQuery: Query;
   private readonly importModuleAliasesQuery: Query;
@@ -88,9 +79,7 @@ export class ElmLsDiagnostics {
   private readonly unusedPortModuleQuery: Query;
   private readonly operatorFunctionsQuery: Query;
   private readonly typeAliasesQuery: Query;
-  private readonly typeAliasUsagesQuery: Query;
   private readonly unionVariantsQuery: Query;
-  private readonly unionVariantUsagesQuery: Query;
   private readonly patternReferencesQuery: Query;
 
   constructor() {
@@ -113,20 +102,6 @@ export class ElmLsDiagnostics {
       `,
     );
 
-    this.exposedValueAndTypeUsagesQuery = this.language.query(
-      `
-      (
-        [
-          (value_expr)
-          (record_base_identifier)
-        ] @value.reference
-      )
-      ((type_ref
-        (upper_case_qid) @type.reference)
-      )
-      `,
-    );
-
     this.moduleImportsQuery = this.language.query(
       `
         (import_clause
@@ -360,20 +335,6 @@ export class ElmLsDiagnostics {
         `,
     );
 
-    this.typeAliasUsagesQuery = this.language.query(
-      `
-        (
-          [
-            (value_expr)
-            (exposed_type)
-          ] @value.reference
-        )
-        ((type_ref
-          (upper_case_qid) @type.reference)
-        )
-        `,
-    );
-
     this.unionVariantsQuery = this.language.query(
       `
         (type_declaration
@@ -385,25 +346,6 @@ export class ElmLsDiagnostics {
         `,
     );
 
-    this.unionVariantUsagesQuery = this.language.query(
-      `
-      (
-        (exposed_type) @exposed.reference
-      )
-      (
-        (value_expr) @value.reference
-      )
-      ((type_ref
-        (upper_case_qid) @type.reference)
-      )
-      ((case_of_branch
-        (pattern
-          (union_pattern
-            (upper_case_qid) @variant.reference)))
-      )
-      `,
-    );
-
     this.patternReferencesQuery = this.language.query(
       `
         (
@@ -671,21 +613,18 @@ export class ElmLsDiagnostics {
         return alias ? alias : node;
       });
 
-    const moduleReferences = this.moduleReferencesQuery
-      .matches(tree.rootNode)
-      .filter(Utils.notUndefined.bind(this))
-      .filter(
-        (match) =>
-          match.captures.length > 0 &&
-          match.captures[0].node.parent?.type !== "import_clause",
-      )
-      .map((match) => match.captures.map((n) => n.node.text).join("."));
-
     // Would need to adjust tree-sitter (use fields) to get a better query
     moduleImports.forEach((moduleImport) => {
-      const references = moduleReferences.filter(
-        (moduleReference) => moduleReference === moduleImport.text,
-      );
+      const references = this.moduleReferencesQuery
+        .matches(tree.rootNode)
+        .filter(Utils.notUndefined.bind(this))
+        .filter(
+          (match) =>
+            match.captures.length > 0 &&
+            match.captures[0].node.parent?.type !== "import_clause",
+        )
+        .map((match) => match.captures.map((n) => n.node.text).join("."))
+        .filter((moduleReference) => moduleReference === moduleImport.text);
 
       const importNode =
         moduleImport.parent?.type === "as_clause"
@@ -713,17 +652,29 @@ export class ElmLsDiagnostics {
       .matches(tree.rootNode)
       .map((match) => match.captures[0].node);
 
-    const allUsages = this.exposedValueAndTypeUsagesQuery
-      .matches(tree.rootNode)
-      .filter(Utils.notUndefined.bind(this));
     exposedValuesAndTypes.forEach((exposedValueOrType) => {
       if (exposedValueOrType.text.endsWith("(..)")) {
         return;
       }
 
-      const references = allUsages.filter(
-        (result) => result.captures[0].node.text === exposedValueOrType.text,
-      );
+      const references = this.language
+        .query(
+          `
+          (
+            [
+              (value_expr)
+              (record_base_identifier)
+            ] @value.reference
+            (#eq? @value.reference "${exposedValueOrType.text}")
+          )
+          ((type_ref
+            (upper_case_qid) @type.reference)
+            (#eq? @type.reference "${exposedValueOrType.text}")
+          )
+          `,
+        )
+        .matches(tree.rootNode)
+        .filter(Utils.notUndefined.bind(this));
 
       if (references.length === 0) {
         diagnostics.push({
@@ -749,21 +700,18 @@ export class ElmLsDiagnostics {
       .matches(tree.rootNode)
       .map((match) => match.captures[0].node);
 
-    const allAliasReferences = this.moduleAliasReferencesQuery
-      .matches(tree.rootNode)
-      .filter(Utils.notUndefined.bind(this))
-      .filter((match) => match.captures.length > 0)
-      .map((match) => match.captures[0].node.text);
-
     moduleAliases.forEach((moduleAlias) => {
       // This case is handled by unused_import
       if (!moduleAlias.parent?.parent?.childForFieldName("exposing")) {
         return;
       }
 
-      const references = allAliasReferences.filter(
-        (moduleReference) => moduleReference === moduleAlias.text,
-      );
+      const references = this.moduleAliasReferencesQuery
+        .matches(tree.rootNode)
+        .filter(Utils.notUndefined.bind(this))
+        .filter((match) => match.captures.length > 0)
+        .map((match) => match.captures[0].node.text)
+        .filter((moduleReference) => moduleReference === moduleAlias.text);
 
       if (references.length === 0 && moduleAlias.parent) {
         diagnostics.push({
@@ -785,8 +733,6 @@ export class ElmLsDiagnostics {
 
     const patternMatches = this.patternsQuery.matches(tree.rootNode);
 
-    const scopeCache = new SyntaxNodeMap<SyntaxNode, QueryResult[]>();
-
     patternMatches
       .filter(Utils.notUndefined.bind(this))
       .map((match) => {
@@ -819,12 +765,9 @@ export class ElmLsDiagnostics {
       })
       .reduce((a, b) => a.concat(b), [])
       .forEach(({ scope, pattern }) => {
-        const references = scopeCache
-          .getOrSet(scope, () =>
-            this.patternReferencesQuery
-              .matches(scope)
-              .filter(Utils.notUndefined.bind(this)),
-          )
+        const references = this.patternReferencesQuery
+          .matches(scope)
+          .filter(Utils.notUndefined.bind(this))
           .filter(
             (result) =>
               result.captures[0].name !== "patternVariable.reference" ||
@@ -1111,14 +1054,26 @@ export class ElmLsDiagnostics {
       .map((match) => match.captures[0].node)
       .filter(Utils.notUndefinedOrNull.bind(this));
 
-    const typeAliasUsages = this.typeAliasUsagesQuery
-      .matches(tree.rootNode)
-      .filter(Utils.notUndefined.bind(this));
-
     typeAliases.forEach((typeAlias) => {
-      const references = typeAliasUsages.filter(
-        (result) => result.captures[0].node.text === typeAlias.text,
-      );
+      const references = this.language
+        .query(
+          `
+          (
+            [
+              (value_expr)
+              (exposed_type)
+            ] @value.reference
+            (#eq? @value.reference "${typeAlias.text}")
+          )
+          ((type_ref
+            (upper_case_qid) @type.reference)
+            (#eq? @type.reference "${typeAlias.text}")
+          )
+          `,
+        )
+        .matches(tree.rootNode)
+        .filter(Utils.notUndefined.bind(this));
+
       if (references.length === 0 && typeAlias.parent) {
         diagnostics.push({
           range: this.getNodeRange(typeAlias.parent),
@@ -1142,18 +1097,32 @@ export class ElmLsDiagnostics {
       .map((match) => [match.captures[1].node, match.captures[0].node])
       .filter(Utils.notUndefinedOrNull.bind(this));
 
-    const unionVariantUsages = this.unionVariantUsagesQuery
-      .matches(tree.rootNode)
-      .filter(Utils.notUndefined.bind(this));
-
     unionVariants.forEach(([unionVariant, typeName]) => {
-      const references = unionVariantUsages.filter(
-        (result) =>
-          result.captures[0].node.text ===
-          (result.captures[0].name === "exposed.reference"
-            ? `${typeName.text}(..)`
-            : unionVariant.text),
-      );
+      const references = this.language
+        .query(
+          `
+          (
+            (exposed_type) @exposed.reference
+            (#eq? @exposed.reference "${typeName.text}(..)")
+          )
+          (
+            (value_expr) @value.reference
+            (#eq? @value.reference "${unionVariant.text}")
+          )
+          ((type_ref
+            (upper_case_qid) @type.reference)
+            (#eq? @type.reference "${unionVariant.text}")
+          )
+          ((case_of_branch
+            (pattern
+              (union_pattern
+                (upper_case_qid) @variant.reference)))
+            (#eq? @variant.reference "${unionVariant.text}")
+          )
+          `,
+        )
+        .matches(tree.rootNode)
+        .filter(Utils.notUndefined.bind(this));
 
       if (references.length === 0 && unionVariant.parent) {
         diagnostics.push({
diff --git a/src/util/types/typeInference.ts b/src/util/types/typeInference.ts
index f842d50..0231338 100644
--- a/src/util/types/typeInference.ts
+++ b/src/util/types/typeInference.ts
@@ -1473,7 +1473,7 @@ export class InferenceScope {
       }
     }
 
-    const type = targetTy.fields[fieldIdentifierText] ?? TUnknown;
+    const type = targetTy.fields[fieldIdentifierText] ?? TVar("a");
     this.expressionTypes.set(expr, type);
     return type;
   }
@@ -1592,6 +1592,7 @@ export class InferenceScope {
               record,
               this.calculateRecordDiff(TRecord(mappedFields), baseType),
             );
+            this.expressionTypes.set(field, type);
           }
         } else if (baseType.nodeType === "MutableRecord") {
           baseType.fields[fieldText] = type;
@@ -2007,7 +2008,6 @@ export class InferenceScope {
         this.bindPattern(field, TUnknown, isParameter);
       });
 
-      this.expressionTypes.set(pattern, type);
       return;
     }
 


Test patch:
diff --git a/test/codeActionProvider.test.ts b/test/codeActionProvider.test.ts
deleted file mode 100644
index f9e25bc..0000000
--- a/test/codeActionProvider.test.ts
+++ /dev/null
@@ -1,190 +0,0 @@
-import { container } from "tsyringe";
-import { CodeAction } from "vscode-languageserver";
-import { URI } from "vscode-uri";
-import { IElmWorkspace } from "../src/elmWorkspace";
-import {
-  CodeActionProvider,
-  convertFromAnalyzerDiagnostic,
-} from "../src/providers";
-import { ICodeActionParams } from "../src/providers/paramsExtensions";
-import { Utils } from "../src/util/utils";
-import { getTargetPositionFromSource } from "./utils/sourceParser";
-import { baseUri, SourceTreeParser } from "./utils/sourceTreeParser";
-
-function codeActionEquals(a: CodeAction, b: CodeAction): boolean {
-  return a.title === b.title;
-}
-
-const basicsSources = `
---@ Basics.elm
-module Basics exposing ((+), (|>), (==), Int, Float, Bool(..), Order(..))
-
-infix left  0 (|>) = apR
-infix non   4 (==) = eq
-infix left  6 (+)  = add
-
-type Int = Int
-
-type Float = Float
-
-type Bool = True | False
-
-add : number -> number -> number
-add =
-  Elm.Kernel.Basics.add
-
-apR : a -> (a -> b) -> b
-apR x f =
-  f x
-
-eq : a -> a -> Bool
-eq =
-  Elm.Kernel.Utils.equal
-
-type Order = LT | EQ | GT
-`;
-
-class MockCodeActionsProvider extends CodeActionProvider {
-  public handleCodeAction(params: ICodeActionParams): CodeAction[] | undefined {
-    return this.onCodeAction(params);
-  }
-}
-
-describe("test codeActionProvider", () => {
-  const treeParser = new SourceTreeParser();
-  let codeActionProvider: MockCodeActionsProvider;
-
-  const debug = process.argv.find((arg) => arg === "--debug");
-
-  async function testCodeAction(
-    source: string,
-    expectedCodeActions: CodeAction[],
-  ) {
-    await treeParser.init();
-
-    if (!codeActionProvider) {
-      codeActionProvider = new MockCodeActionsProvider();
-    }
-
-    const result = getTargetPositionFromSource(source);
-
-    if (!result) {
-      throw new Error("Getting sources failed");
-    }
-
-    const testUri = URI.file(baseUri + "Test.elm").toString();
-
-    const program = await treeParser.getProgram(result.sources);
-    const sourceFile = program.getForest().getByUri(testUri);
-
-    if (!sourceFile) throw new Error("Getting tree failed");
-
-    const workspaces = container.resolve<IElmWorkspace[]>("ElmWorkspaces");
-    workspaces.splice(0, workspaces.length);
-    workspaces.push(program);
-
-    const range = { start: result.position, end: result.position };
-    const codeActions =
-      codeActionProvider.handleCodeAction({
-        program,
-        sourceFile,
-        range,
-        textDocument: { uri: testUri },
-        context: {
-          diagnostics: [
-            ...program.getSyntacticDiagnostics(sourceFile),
-            ...program.getSemanticDiagnostics(sourceFile),
-            ...program.getSuggestionDiagnostics(sourceFile),
-          ]
-            .filter((diag) => Utils.rangeOverlaps(diag.range, range))
-            .map(convertFromAnalyzerDiagnostic),
-        },
-      }) ?? [];
-
-    const codeActionsExist = expectedCodeActions.every((codeAction) =>
-      codeActions.find((c) => codeActionEquals(codeAction, c)),
-    );
-
-    if (debug && !codeActionsExist) {
-      console.log(
-        `Expecting ${JSON.stringify(expectedCodeActions)}, got ${JSON.stringify(
-          codeActions,
-        )}`,
-      );
-    }
-
-    expect(codeActionsExist).toBeTruthy();
-  }
-
-  test("add import of value", async () => {
-    const source = `
---@ Test.elm
-module Test exposing (..)
-
-func = foo
-      --^
-
---@ App.elm
-module App exposing (foo)
-
-foo = ""
-`;
-    await testCodeAction(basicsSources + source, [
-      { title: `Import 'foo' from module "App"` },
-    ]);
-  });
-
-  test("add import of qualified value", async () => {
-    const source = `
---@ Test.elm
-module Test exposing (..)
-
-func = App.foo
-          --^
-
---@ App.elm
-module App exposing (foo)
-
-foo = ""
-`;
-    await testCodeAction(basicsSources + source, [
-      { title: `Import module "App"` },
-    ]);
-
-    const source2 = `
---@ Test.elm
-module Test exposing (..)
-
-func = App.foo
-      --^
-
---@ App.elm
-module App exposing (foo)
-
-foo = ""
-`;
-    await testCodeAction(basicsSources + source2, [
-      { title: `Import module "App"` },
-    ]);
-  });
-
-  test("add all missing imports", async () => {
-    const source = `
---@ Test.elm
-module Test exposing (..)
-
-func = foo + bar
-      --^
-
---@ App.elm
-module App exposing (foo, bar)
-
-foo = ""
-
-bar = ""
-`;
-    await testCodeAction(basicsSources + source, [
-      { title: `Add all missing imports` },
-    ]);
-  });
-});
diff --git a/test/codeActionTests/addMissingRecordFieldCodeAction.test.ts b/test/codeActionTests/addMissingRecordFieldCodeAction.test.ts
new file mode 100644
index 0000000..657f641
--- /dev/null
+++ b/test/codeActionTests/addMissingRecordFieldCodeAction.test.ts
@@ -0,0 +1,97 @@
+import { testCodeAction } from "./codeActionTestBase";
+
+describe("add missing record field code action", () => {
+  it("should work with type alias on a single line", async () => {
+    const source = `
+		--@ Test.elm
+module Test exposing (..)
+
+type alias Model = { field : Int }
+
+func : Model -> Float
+func model =
+		model.newProp
+			     --^
+		`;
+
+    const expectedSource = `
+		--@ Test.elm
+module Test exposing (..)
+
+type alias Model = { field : Int, newProp : Float }
+
+func : Model -> Float
+func model =
+		model.newProp
+		`;
+
+    await testCodeAction(
+      source,
+      [{ title: "Create missing record field" }],
+      expectedSource,
+    );
+  });
+
+  it("should work with type alias on a multiple lines", async () => {
+    const source = `
+		--@ Test.elm
+module Test exposing (..)
+
+type alias Model = 
+  { field : Int
+	}
+
+func : Model -> { prop : Int }
+func model =
+		model.newProp
+			     --^
+		`;
+
+    const expectedSource = `
+		--@ Test.elm
+module Test exposing (..)
+
+type alias Model = 
+  { field : Int
+  , newProp : { prop : Int }
+	}
+
+func : Model -> { prop : Int }
+func model =
+		model.newProp
+		`;
+
+    await testCodeAction(
+      source,
+      [{ title: "Create missing record field" }],
+      expectedSource,
+    );
+  });
+
+  it("should work with record update expr", async () => {
+    const source = `
+		--@ Test.elm
+module Test exposing (..)
+
+func : { field : Int } -> a
+func model =
+		{ model | newProp = "" } 
+			         --^
+		`;
+
+    const expectedSource = `
+		--@ Test.elm
+module Test exposing (..)
+
+func : { field : Int, newProp : String } -> a
+func model =
+		{ model | newProp = "" } 
+		`;
+
+    await testCodeAction(
+      source,
+      [{ title: "Create missing record field" }],
+      expectedSource,
+    );
+  });
+});
diff --git a/test/codeActionTests/codeActionTestBase.ts b/test/codeActionTests/codeActionTestBase.ts
new file mode 100644
index 0000000..69569c1
--- /dev/null
+++ b/test/codeActionTests/codeActionTestBase.ts
@@ -0,0 +1,139 @@
+import { container } from "tsyringe";
+import { CodeAction } from "vscode-languageserver";
+import { URI } from "vscode-uri";
+import { IElmWorkspace } from "../../src/elmWorkspace";
+import {
+  CodeActionProvider,
+  convertFromAnalyzerDiagnostic,
+} from "../../src/providers";
+import { ICodeActionParams } from "../../src/providers/paramsExtensions";
+import { Utils } from "../../src/util/utils";
+import {
+  getTargetPositionFromSource,
+  getSourceFiles,
+} from "../utils/sourceParser";
+import {
+  SourceTreeParser,
+  trimTrailingWhitespace,
+  baseUri,
+  applyEditsToSource,
+  stripCommentLines,
+} from "../utils/sourceTreeParser";
+
+function codeActionEquals(a: CodeAction, b: CodeAction): boolean {
+  return a.title === b.title;
+}
+
+const basicsSources = `
+--@ Basics.elm
+module Basics exposing ((+), (|>), (==), Int, Float, Bool(..), Order(..))
+
+infix left  0 (|>) = apR
+infix non   4 (==) = eq
+infix left  6 (+)  = add
+
+type Int = Int
+
+type Float = Float
+
+type Bool = True | False
+
+add : number -> number -> number
+add =
+  Elm.Kernel.Basics.add
+
+apR : a -> (a -> b) -> b
+apR x f =
+  f x
+
+eq : a -> a -> Bool
+eq =
+  Elm.Kernel.Utils.equal
+
+type Order = LT | EQ | GT
+`;
+
+class MockCodeActionsProvider extends CodeActionProvider {
+  public handleCodeAction(params: ICodeActionParams): CodeAction[] | undefined {
+    return this.onCodeAction(params);
+  }
+}
+
+export async function testCodeAction(
+  source: string,
+  expectedCodeActions: CodeAction[],
+  expectedResultAfterEdits?: string,
+): Promise<void> {
+  const treeParser = new SourceTreeParser();
+  await treeParser.init();
+  const codeActionProvider = new MockCodeActionsProvider();
+
+  const debug = process.argv.find((arg) => arg === "--debug");
+
+  const result = getTargetPositionFromSource(
+    basicsSources + trimTrailingWhitespace(source),
+  );
+
+  if (!result) {
+    throw new Error("Getting sources failed");
+  }
+
+  const testUri = URI.file(baseUri + "Test.elm").toString();
+
+  const program = await treeParser.getProgram(result.sources);
+  const sourceFile = program.getForest().getByUri(testUri);
+
+  if (!sourceFile) throw new Error("Getting tree failed");
+
+  const workspaces = container.resolve<IElmWorkspace[]>("ElmWorkspaces");
+  workspaces.splice(0, workspaces.length);
+  workspaces.push(program);
+
+  const range = { start: result.position, end: result.position };
+  const codeActions =
+    codeActionProvider.handleCodeAction({
+      program,
+      sourceFile,
+      range,
+      textDocument: { uri: testUri },
+      context: {
+        diagnostics: [
+          ...program.getSyntacticDiagnostics(sourceFile),
+          ...program.getSemanticDiagnostics(sourceFile),
+          ...program.getSuggestionDiagnostics(sourceFile),
+        ]
+          .filter((diag) => Utils.rangeOverlaps(diag.range, range))
+          .map(convertFromAnalyzerDiagnostic),
+      },
+    }) ?? [];
+
+  const codeActionsExist = expectedCodeActions.every((codeAction) =>
+    codeActions.find((c) => codeActionEquals(codeAction, c)),
+  );
+
+  if (debug && !codeActionsExist) {
+    console.log(
+      `Expecting ${JSON.stringify(expectedCodeActions)}, got ${JSON.stringify(
+        codeActions,
+      )}`,
+    );
+  }
+
+  expect(codeActionsExist).toBeTruthy();
+
+  if (expectedResultAfterEdits) {
+    const expectedSources = getSourceFiles(
+      trimTrailingWhitespace(expectedResultAfterEdits),
+    );
+
+    Object.entries(expectedSources).forEach(([uri, source]) => {
+      expect(
+        applyEditsToSource(
+          stripCommentLines(result.sources[uri]),
+          codeActions[0].edit?.changes![URI.file(baseUri + uri).toString()] ??
+            [],
+        ),
+      ).toEqual(source);
+    });
+  }
+}
diff --git a/test/codeActionTests/importCodeAction.test.ts b/test/codeActionTests/importCodeAction.test.ts
new file mode 100644
index 0000000..0c1edbc
--- /dev/null
+++ b/test/codeActionTests/importCodeAction.test.ts
@@ -0,0 +1,67 @@
+import { testCodeAction } from "./codeActionTestBase";
+
+describe("import code actions", () => {
+  test("add import of value", async () => {
+    const source = `
+--@ Test.elm
+module Test exposing (..)
+
+func = foo
+      --^
+
+--@ App.elm
+module App exposing (foo)
+
+foo = ""
+`;
+    await testCodeAction(source, [{ title: `Import 'foo' from module "App"` }]);
+  });
+
+  test("add import of qualified value", async () => {
+    const source = `
+--@ Test.elm
+module Test exposing (..)
+
+func = App.foo
+          --^
+
+--@ App.elm
+module App exposing (foo)
+
+foo = ""
+`;
+    await testCodeAction(source, [{ title: `Import module "App"` }]);
+
+    const source2 = `
+--@ Test.elm
+module Test exposing (..)
+
+func = App.foo
+      --^
+
+--@ App.elm
+module App exposing (foo)
+
+foo = ""
+`;
+    await testCodeAction(source2, [{ title: `Import module "App"` }]);
+  });
+
+  test("add all missing imports", async () => {
+    const source = `
+--@ Test.elm
+module Test exposing (..)
+
+func = foo + bar
+      --^
+
+--@ App.elm
+module App exposing (foo, bar)
+
+foo = ""
+
+bar = ""
+`;
+    await testCodeAction(source, [{ title: `Add all missing imports` }]);
+  });
+});
diff --git a/test/completionProvider.test.ts b/test/completionProvider.test.ts
index 7a50696..faeebdf 100644
--- a/test/completionProvider.test.ts
+++ b/test/completionProvider.test.ts
@@ -1410,55 +1410,4 @@ func2 =
 
     await testCompletions(source2, ["func", "test"], "partialMatch");
   });
-
-  it("Record completions in record patterns", async () => {
-    const source = `
---@ Test.elm
-module Test exposing (..)
-
-func : { prop1: String, prop2: Int } -> String
-func {{-caret-}} =
-    ""
-`;
-
-    await testCompletions(source, ["prop1", "prop2"], "exactMatch");
-
-    const source2 = `
---@ Test.elm
-module Test exposing (..)
-
-func : { prop1: String, prop2: Int } -> String
-func {p{-caret-}} =
-    ""
-`;
-
-    await testCompletions(source2, ["prop1", "prop2"], "exactMatch");
-
-    const source3 = `
---@ Test.elm
-module Test exposing (..)
-
-func =
-    let
-      {{-caret-}} = { prop1 = "", prop2 = 2 }
-    in
-    ""
-`;
-
-    await testCompletions(source3, ["prop1", "prop2"], "exactMatch");
-  });
-
-  it("Case branch patterns should have completions", async () => {
-    const source = `
---@ Test.elm
-module Test exposing (..)
-
-func model =
-    case model of 
-      Just { details } ->
-        d{-caret-}
-`;
-
-    await testCompletions(source, ["details"]);
-  });
 });