Fix Substance AST queries test

penrose · May 11, 2022 · a82da0f · a82da0f
1 parent 9eb147d
commit a82da0f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/packages/core/src/analysis/SubstanceAnalysis.test.ts b/packages/core/src/analysis/SubstanceAnalysis.test.ts
@@ -6,8 +6,8 @@ import {
 } from "compiler/Substance";
 import { dummyIdentifier } from "engine/EngineUtils";
 import { intersectionWith } from "lodash";
-import { similarMappings, similarNodes } from "synthesis/Search";
-import { A, ASTNode } from "types/ast";
+import { similarMappings, similarNodes, SubNode } from "synthesis/Search";
+import { A } from "types/ast";
 import { Env } from "types/domain";
 import { SubProg, SubStmt } from "types/substance";
 import {
@@ -33,8 +33,8 @@ const compile = (src: string): SubProg<A> =>
 
 describe("Substance AST queries", () => {
   test("Similar AST nodes", () => {
-    let node1: ASTNode<A>;
-    let node2: ASTNode<A>;
+    let node1: SubNode<A>;
+    let node2: SubNode<A>;
     node1 = compile("Set A");
     node2 = compile("Set B");
     expect(similarNodes(node1, node2)).toBe(true);

diff --git a/packages/core/src/synthesis/Search.ts b/packages/core/src/synthesis/Search.ts
@@ -116,7 +116,12 @@ export const diffSubStmts = (
   return exactDiffs.map((d) => toStmtDiff(d, leftSorted));
 };
 
-type SubNode<T> = LabelOption<T> | SubExpr<T> | SubStmt<T> | Type<T>;
+export type SubNode<T> =
+  | LabelOption<T>
+  | SubExpr<T>
+  | SubProg<T>
+  | SubStmt<T>
+  | Type<T>;
 
 const children = <T>(node: SubNode<T>): SubNode<T>[] => {
   switch (node.tag) {
@@ -158,6 +163,9 @@ const children = <T>(node: SubNode<T>): SubNode<T>[] => {
     case "NoLabel": {
       return [...node.args];
     }
+    case "SubProg": {
+      return [...node.statements];
+    }
     case "TypeVar": {
       return [node.name];
     }