From 6bdbf14ecb4af8d48227dc6c50820d509b706c76 Mon Sep 17 00:00:00 2001
From: Nilstrieb <48135649+Nilstrieb@users.noreply.github.com>
Date: Mon, 31 Jul 2023 13:13:30 +0200
Subject: [PATCH] Fix folders of items wrt the itemid map

---
 src/ast.ts     | 22 +++++++++++++++++++++-
 src/parser.ts  | 10 ++++++++--
 src/resolve.ts |  5 ++++-
 src/typeck.ts  |  5 ++++-
 4 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/src/ast.ts b/src/ast.ts
index ba6d629..831cabd 100644
--- a/src/ast.ts
+++ b/src/ast.ts
@@ -447,14 +447,29 @@ export type TypeckResults = {
 export type FoldFn<T> = (value: T) => T;
 
 export type Folder = {
+  ast: () => Ast;
+  /**
+   * This should not be overridden.
+   */
   item: FoldFn<Item>;
+  itemInner: FoldFn<Item>;
   expr: FoldFn<Expr>;
   ident: FoldFn<Identifier>;
   type: FoldFn<Type>;
 };
 
+const ITEM_DEFAULT: symbol = Symbol("item must not be overriden");
+
 export const DEFAULT_FOLDER: Folder = {
+  ast() {
+    throw new Error("folders need to implement `ast`");
+  },
   item(item) {
+    const newItem = this.itemInner(item);
+    this.ast().itemsById.set(item.id, item);
+    return newItem;
+  },
+  itemInner(item) {
     return superFoldItem(item, this);
   },
   expr(expr) {
@@ -467,8 +482,13 @@ export const DEFAULT_FOLDER: Folder = {
     return superFoldType(type, this);
   },
 };
+(DEFAULT_FOLDER.item as any)[ITEM_DEFAULT] = ITEM_DEFAULT;
 
 export function foldAst(ast: Ast, folder: Folder): Ast {
+  if ((folder.item as any)[ITEM_DEFAULT] !== ITEM_DEFAULT) {
+    throw new Error("must not override `item` on folders");
+  }
+
   return {
     rootItems: ast.rootItems.map((item) => folder.item(item)),
     itemsById: ast.itemsById,
@@ -528,7 +548,7 @@ export function superFoldItem(item: Item, folder: Folder): Item {
     }
     case "mod": {
       let kind: ModItemKind;
-      const { modKind: itemKind } = item.node;    
+      const { modKind: itemKind } = item.node;
       switch (itemKind.kind) {
         case "inline":
           kind = {
diff --git a/src/parser.ts b/src/parser.ts
index 9115141..3bf1834 100644
--- a/src/parser.ts
+++ b/src/parser.ts
@@ -658,7 +658,10 @@ function validateAst(ast: Ast) {
 
   const validator: Folder = {
     ...DEFAULT_FOLDER,
-    item(item: Item): Item {
+    ast() {
+      return ast;
+    },
+    itemInner(item: Item): Item {
       if (seenItemIds.has(item.id)) {
         throw new Error(`duplicate item id: ${item.id} for ${item.node.name}`);
       }
@@ -719,7 +722,10 @@ function assignIds(rootItems: Item[]): Ast {
 
   const assigner: Folder = {
     ...DEFAULT_FOLDER,
-    item(item: Item): Item {
+    ast() {
+      return ast;
+    },
+    itemInner(item: Item): Item {
       const id = itemId.next();
       ast.itemsById.set(id, item);
       return { ...superFoldItem(item, this), id };
diff --git a/src/resolve.ts b/src/resolve.ts
index c244f98..20377b5 100644
--- a/src/resolve.ts
+++ b/src/resolve.ts
@@ -113,7 +113,10 @@ function resolveModule(cx: Context, contents: Item[]): Item[] {
 
   const resolver: Folder = {
     ...DEFAULT_FOLDER,
-    item(item) {
+    ast() {
+      return cx.ast;
+    },
+    itemInner(item) {
       switch (item.kind) {
         case "function": {
           const params = item.node.params.map(({ name, span, type }) => ({
diff --git a/src/typeck.ts b/src/typeck.ts
index 05ee403..df57f0c 100644
--- a/src/typeck.ts
+++ b/src/typeck.ts
@@ -189,7 +189,10 @@ export function typeck(ast: Ast): Ast {
 
   const checker: Folder = {
     ...DEFAULT_FOLDER,
-    item(item) {
+    ast() {
+      return ast;
+    },
+    itemInner(item) {
       switch (item.kind) {
         case "function": {
           const fnTy = typeOfItem(item.id, item.span) as TyFn;