From 068931e9ed65b568e28581b491141cba0191f8c0 Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Sat, 23 Nov 2024 09:32:52 -0800
Subject: [PATCH] fix: proper semantic highlighting and formatting for type
 declarations

---
 packages/schema/src/language-server/zmodel-formatter.ts | 8 ++++----
 packages/schema/src/language-server/zmodel-semantic.ts  | 5 ++++-
 2 files changed, 8 insertions(+), 5 deletions(-)

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) ||