From b7c6c87a2231a8dae1f0735a027926e2608ac2c8 Mon Sep 17 00:00:00 2001
From: Yiming <yiming@whimslab.io>
Date: Sat, 23 Nov 2024 10:00:07 -0800
Subject: [PATCH] fix: proper semantic highlighting and formatting for type
 declarations (#1883)

---
 packages/ide/jetbrains/CHANGELOG.md                     | 6 ++++++
 packages/schema/src/language-server/zmodel-formatter.ts | 8 ++++----
 packages/schema/src/language-server/zmodel-semantic.ts  | 5 ++++-
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/packages/ide/jetbrains/CHANGELOG.md b/packages/ide/jetbrains/CHANGELOG.md
index 79da788cc..48cdd8b69 100644
--- a/packages/ide/jetbrains/CHANGELOG.md
+++ b/packages/ide/jetbrains/CHANGELOG.md
@@ -2,6 +2,12 @@
 
 ## [Unreleased]
 
+### Fixed
+
+-   Proper semantic highlighting and formatting for type declarations.
+
+## 2.9.0
+
 ### Added
 
 -   Support for using `@@validate` attribute inside type declarations.
diff --git a/packages/schema/src/language-server/zmodel-formatter.ts b/packages/schema/src/language-server/zmodel-formatter.ts
index 3a17520b0..92296176a 100644
--- a/packages/schema/src/language-server/zmodel-formatter.ts
+++ b/packages/schema/src/language-server/zmodel-formatter.ts
@@ -26,8 +26,8 @@ export class ZModelFormatter extends AbstractFormatter {
     protected format(node: AstNode): void {
         const formatter = this.getNodeFormatter(node);
 
-        if (ast.isDataModelField(node)) {
-            if (this.isPrismaStyle && ast.isDataModel(node.$container)) {
+        if (ast.isDataModelField(node) || ast.isTypeDefField(node)) {
+            if (this.isPrismaStyle && (ast.isDataModel(node.$container) || ast.isTypeDef(node.$container))) {
                 const dataModel = node.$container;
 
                 const compareFn = (a: number, b: number) => b - a;
@@ -104,14 +104,14 @@ export class ZModelFormatter extends AbstractFormatter {
         this.isPrismaStyle = isPrismaStyle;
     }
 
-    private getFieldTypeLength(field: ast.DataModelField) {
+    private getFieldTypeLength(field: ast.DataModelField | ast.TypeDefField) {
         let length: number;
 
         if (field.type.type) {
             length = field.type.type.length;
         } else if (field.type.reference) {
             length = field.type.reference.$refText.length;
-        } else if (field.type.unsupported) {
+        } else if (ast.isDataModelField(field) && field.type.unsupported) {
             const name = `Unsupported("${field.type.unsupported.value.value}")`;
             length = name.length;
         } else {
diff --git a/packages/schema/src/language-server/zmodel-semantic.ts b/packages/schema/src/language-server/zmodel-semantic.ts
index 18e2e5819..2e24cdb7c 100644
--- a/packages/schema/src/language-server/zmodel-semantic.ts
+++ b/packages/schema/src/language-server/zmodel-semantic.ts
@@ -18,6 +18,8 @@ import {
     isPlugin,
     isPluginField,
     isReferenceExpr,
+    isTypeDef,
+    isTypeDefField,
 } from '@zenstackhq/language/ast';
 import { AbstractSemanticTokenProvider, AstNode, SemanticTokenAcceptor } from 'langium';
 import { SemanticTokenTypes } from 'vscode-languageserver';
@@ -36,7 +38,7 @@ export class ZModelSemanticTokenProvider extends AbstractSemanticTokenProvider {
                 property: 'superTypes',
                 type: SemanticTokenTypes.type,
             });
-        } else if (isDataSource(node) || isGeneratorDecl(node) || isPlugin(node) || isEnum(node)) {
+        } else if (isDataSource(node) || isGeneratorDecl(node) || isPlugin(node) || isEnum(node) || isTypeDef(node)) {
             acceptor({
                 node,
                 property: 'name',
@@ -44,6 +46,7 @@ export class ZModelSemanticTokenProvider extends AbstractSemanticTokenProvider {
             });
         } else if (
             isDataModelField(node) ||
+            isTypeDefField(node) ||
             isConfigField(node) ||
             isAttributeArg(node) ||
             isPluginField(node) ||