feat: support permission type annotations (use typechecking) (#21)

Parse the optional return-type annotation on permissions
(`permission view: user | group = ...`) introduced by the SpiceDB
`use typechecking` flag. Previously any annotated permission failed to
parse, which — in editor tooling — disabled go-to-definition and
semantic highlighting for the entire schema.

- dsl.ts: add an optional `: type | type` annotation to the permission
  grammar, restricted to plain type identifiers to match the SpiceDB
  parser. Expose it as `ParsedPermission.annotatedTypes` (a TypeExpr of
  TypeRefs) and walk it in findReferenceNode and mapParseNodes so
  go-to-definition resolves on annotation types.
- resolution.ts: include annotation type refs in resolvedReferences()
  so they resolve and highlight like relation types.
- Tests for the AST, the error cases mirrored from the SpiceDB parser
  fixtures, findReferenceNode, and the resolver.

Author: josephschorr

src/dsl.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "vitest";
-import { parseSchema, stringLiteral } from "./dsl";
+import { findReferenceNode, parse, parseSchema, stringLiteral } from "./dsl";
 
 // Takes an expression as an argument and throws if that assertion fails.
 // This primarily exists to provide typescript narrowing in a statement,
@@ -495,6 +495,113 @@ definition user {}
     });
   });
 
+  describe("permission type annotations", () => {
+    it("parses a permission with a single type annotation", () => {
+      const schema = `definition foo {
+          permission view: user = viewer
+        }`;
+      const parsed = parseSchema(schema);
+
+      const definition = parsed?.definitions[0];
+      assert(definition);
+      assert(definition.kind === "objectDef");
+
+      const permission = definition.permissions[0];
+      assert(permission);
+      expect(permission.name).toEqual("view");
+
+      // The compute expression is still parsed normally.
+      assert(permission.expr.kind === "relationref");
+      expect(permission.expr.relationName).toEqual("viewer");
+
+      // The annotation is captured as a type expression.
+      assert(permission.annotatedTypes);
+      expect(permission.annotatedTypes.types.map((t) => t.path)).toEqual([
+        "user",
+      ]);
+      const annoType = permission.annotatedTypes.types[0];
+      assert(annoType);
+      expect(annoType.kind).toEqual("typeref");
+      expect(annoType.relationName).toBeUndefined();
+      expect(annoType.wildcard).toEqual(false);
+    });
+
+    it("parses a permission with multiple piped type annotations", () => {
+      const schema = `definition foo {
+          permission view: user | group | team = viewer
+        }`;
+      const parsed = parseSchema(schema);
+
+      const definition = parsed?.definitions[0];
+      assert(definition);
+      assert(definition.kind === "objectDef");
+
+      const permission = definition.permissions[0];
+      assert(permission);
+      assert(permission.annotatedTypes);
+      expect(permission.annotatedTypes.types.map((t) => t.path)).toEqual([
+        "user",
+        "group",
+        "team",
+      ]);
+    });
+
+    it("leaves annotatedTypes undefined when there is no annotation", () => {
+      const schema = `definition foo {
+          permission view = viewer
+        }`;
+      const parsed = parseSchema(schema);
+
+      const definition = parsed?.definitions[0];
+      assert(definition);
+      assert(definition.kind === "objectDef");
+
+      const permission = definition.permissions[0];
+      assert(permission);
+      expect(permission.annotatedTypes).toBeUndefined();
+    });
+
+    it("rejects a double colon", () => {
+      const schema = `definition foo {
+          permission view:: user = viewer
+        }`;
+      expect(parseSchema(schema)).toBeUndefined();
+    });
+
+    it("rejects a pipe with no preceding type", () => {
+      const schema = `definition foo {
+          permission view: | user = viewer
+        }`;
+      expect(parseSchema(schema)).toBeUndefined();
+    });
+
+    it("rejects a trailing pipe with no following type", () => {
+      const schema = `definition foo {
+          permission view: user | = viewer
+        }`;
+      expect(parseSchema(schema)).toBeUndefined();
+    });
+
+    it("finds the annotation type as a reference node for go-to-definition", () => {
+      const schema = `definition user {}
+definition foo {
+	permission view: user = viewer
+}`;
+      const result = parse(schema);
+      assert(!result.error);
+      assert(result.schema);
+
+      // "\tpermission view: user = viewer" is line 3 (1-indexed).
+      const line = "\tpermission view: user = viewer";
+      const col = line.indexOf("user") + 2; // 1-indexed, inside "user"
+      const found = findReferenceNode(result.schema, 3, col);
+      assert(found);
+      assert(found.node);
+      assert(found.node.kind === "typeref");
+      expect(found.node.path).toEqual("user");
+    });
+  });
+
   describe("partial syntax", () => {
     it("parses a basic partial", () => {
       const schema = `partial thing {

src/dsl.ts

@@ -320,6 +320,12 @@ export type ParsedBinaryExpression = {
 export type ParsedPermission = {
   kind: "permission";
   name: string;
+  /**
+   * annotatedTypes is the optional `use typechecking` return-type annotation on the
+   * permission (e.g. `permission view: user | group = ...`), or undefined if none was
+   * written. Each annotated type is a plain type reference to a definition.
+   */
+  annotatedTypes: TypeExpr | undefined;
   expr: ParsedExpression;
   range: TextRange;
 };
@@ -636,20 +642,59 @@ const tableParser: Parser<ParsedExpression> = table.reduce(
 
 const expr = tableParser.trim(whitespace);
 
+// Permission type annotation (the `use typechecking` return-type annotation).
+// Unlike a relation's allowed types, an annotation is a pipe-separated list of plain
+// type identifiers -- no wildcards, subrelations, caveats, or expiration -- matching
+// the SpiceDB grammar `permission foo: user | group = ...`.
+const annotationTypeRef: Parser<TypeRef> = seqMap(
+  index,
+  identifier,
+  index,
+  function (startIndex, name, endIndex) {
+    return {
+      kind: "typeref",
+      path: name,
+      relationName: undefined,
+      wildcard: false,
+      withCaveat: undefined,
+      withExpiration: undefined,
+      range: { startIndex: startIndex, endIndex: endIndex },
+    };
+  },
+);
+
+const pipedAnnotationTypeRef = pipe.then(annotationTypeRef);
+
+const permissionTypeAnnotation: Parser<TypeExpr> = seqMap(
+  index,
+  seq(annotationTypeRef, pipedAnnotationTypeRef.atLeast(0)),
+  index,
+  function (startIndex, data, endIndex) {
+    const remaining = data[1];
+    return {
+      kind: "typeexpr",
+      types: [data[0], ...remaining],
+      range: { startIndex: startIndex, endIndex: endIndex },
+    };
+  },
+);
+
 // Definitions members.
 const permission: Parser<ParsedPermission> = seqMap(
   index,
   seq(
     lexeme(string("permission")),
     identifier,
+    colon.then(permissionTypeAnnotation).atMost(1),
     equal.then(expr).skip(terminator.atMost(1)),
   ),
   index,
   function (startIndex, data, endIndex) {
     return {
       kind: "permission",
       name: data[1],
-      expr: data[2],
+      annotatedTypes: data[2][0],
+      expr: data[3],
       range: { startIndex: startIndex, endIndex: endIndex },
     };
   },
@@ -918,6 +963,17 @@ function findReferenceNodeInPermission(
   lineNumber: number,
   columnPosition: number,
 ): ParsedRelationRefExpression | TypeRef | undefined {
+  // A `use typechecking` annotation type (e.g. the `user` in `permission view: user = ...`)
+  // is a type reference, so resolve it like a relation's allowed type for go-to-definition.
+  if (permission.annotatedTypes) {
+    const annotatedType = permission.annotatedTypes.types.find(
+      (typeRef: TypeRef) => rangeContains(typeRef, lineNumber, columnPosition),
+    );
+    if (annotatedType) {
+      return annotatedType;
+    }
+  }
+
   const found = flatMapExpression(permission.expr, (expr: ParsedExpression) => {
     if (!rangeContains(expr, lineNumber, columnPosition)) {
       return undefined;
@@ -989,6 +1045,9 @@ const mapParseNodes =
         break;
 
       case "permission":
+        if (node.annotatedTypes) {
+          mapParseNodes(mapper)(node.annotatedTypes);
+        }
         flatMapExpression(node.expr, mapper);
         break;
 

src/resolution.test.ts

@@ -0,0 +1,62 @@
+import { describe, expect, it } from "vitest";
+import { parse } from "./dsl";
+import { Resolver } from "./resolution";
+
+// Takes an expression as an argument and throws if that assertion fails, providing
+// typescript narrowing in a statement.
+function assert(val: unknown, msg = "Assertion failed"): asserts val {
+  if (!val) throw new Error(msg);
+}
+
+describe("resolution", () => {
+  it("resolves a permission type annotation as a type reference", () => {
+    const schema = `definition user {}
+definition document {
+  relation viewer: user
+  permission view: user = viewer
+}`;
+    const result = parse(schema);
+    assert(!result.error);
+    assert(result.schema);
+
+    const resolver = new Resolver(result.schema);
+    const refs = resolver.resolvedReferences();
+
+    // Two type references: the relation's `user` and the annotation's `user`.
+    const typeRefs = refs.filter((r) => r.kind === "type");
+    expect(typeRefs).toHaveLength(2);
+    for (const ref of typeRefs) {
+      expect(ref.reference.path).toEqual("user");
+      assert(ref.referencedTypeAndRelation);
+      expect(ref.referencedTypeAndRelation.definition?.name).toEqual("user");
+    }
+
+    // The compute expression reference `viewer` still resolves to the relation.
+    const exprRefs = refs.filter((r) => r.kind === "expression");
+    expect(exprRefs).toHaveLength(1);
+    const exprRef = exprRefs[0];
+    assert(exprRef);
+    expect(exprRef.reference.relationName).toEqual("viewer");
+    expect(exprRef.resolvedRelationOrPermission?.kind).toEqual("relation");
+  });
+
+  it("marks an unknown annotation type as unresolved", () => {
+    const schema = `definition document {
+  permission view: nonexistent = editor
+}`;
+    const result = parse(schema);
+    assert(!result.error);
+    assert(result.schema);
+
+    const resolver = new Resolver(result.schema);
+    const typeRefs = resolver
+      .resolvedReferences()
+      .filter((r) => r.kind === "type");
+    expect(typeRefs).toHaveLength(1);
+
+    const ref = typeRefs[0];
+    assert(ref);
+    expect(ref.reference.path).toEqual("nonexistent");
+    expect(ref.referencedTypeAndRelation).toBeUndefined();
+  });
+});

src/resolution.ts

@@ -205,15 +205,28 @@ export class Resolver {
         return [];
       }
 
-      return def.relations.flatMap((rel: ParsedRelation) => {
-        return rel.allowedTypes.types.map((typeRef: TypeRef) => {
-          return {
-            kind: "type",
-            reference: typeRef,
-            referencedTypeAndRelation: this.resolveTypeReference(typeRef),
-          };
-        });
-      });
+      const asTypeReference = (typeRef: TypeRef): ResolvedTypeReference => {
+        return {
+          kind: "type",
+          reference: typeRef,
+          referencedTypeAndRelation: this.resolveTypeReference(typeRef),
+        };
+      };
+
+      const relationRefs = def.relations.flatMap((rel: ParsedRelation) =>
+        rel.allowedTypes.types.map(asTypeReference),
+      );
+
+      // Permission `use typechecking` annotations reference types too, so resolve
+      // them alongside relation types for highlighting and hover.
+      const annotationRefs = def.permissions.flatMap(
+        (perm: ParsedPermission) =>
+          perm.annotatedTypes
+            ? perm.annotatedTypes.types.map(asTypeReference)
+            : [],
+      );
+
+      return [...relationRefs, ...annotationRefs];
     });
   }