Merge pull request #9 from distributed-lab/feature/arrays-support

Add resolver for arrays in the component main

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@distributedlab/circom-parser",
   "description": "Circom circuit parser built with ANTLR4",
-  "version": "0.2.1",
+  "version": "0.2.2",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "files": [

src/ExtendedCircomParser.ts

@@ -12,16 +12,24 @@ export class ExtendedCircomParser extends CircomParser {
   parserErrorListener: ErrorListener<Token>;
   lexerErrorListener: ErrorListener<number>;
 
-  constructor(tokens: antlr4.CommonTokenStream, lexer: CircomLexer) {
+  fileIdentifier: string;
+
+  constructor(
+    fileIdentifier: string,
+    tokens: antlr4.CommonTokenStream,
+    lexer: CircomLexer,
+  ) {
     super(tokens);
 
     this.lexer = lexer;
-    this.lexerErrorListener = new ErrorListener();
-    this.parserErrorListener = new ErrorListener();
+    this.lexerErrorListener = new ErrorListener(fileIdentifier);
+    this.parserErrorListener = new ErrorListener(fileIdentifier);
 
     this.initErrorListeners();
 
     this.buildParseTrees = true;
+
+    this.fileIdentifier = fileIdentifier;
   }
 
   circuit() {
@@ -42,11 +50,11 @@ export class ExtendedCircomParser extends CircomParser {
   }
 
   initErrorListeners() {
-    this.parserErrorListener = new ErrorListener();
+    this.parserErrorListener = new ErrorListener(this.fileIdentifier);
     this.removeErrorListeners();
     this.addErrorListener(this.parserErrorListener);
 
-    this.lexerErrorListener = new ErrorListener();
+    this.lexerErrorListener = new ErrorListener(this.fileIdentifier);
     this.lexer.removeErrorListeners();
     this.lexer.addErrorListener(this.lexerErrorListener);
   }

src/ExtendedCircomVisitor.ts

@@ -15,7 +15,7 @@ export class ExtendedCircomVisitor<Result> extends CircomVisitor<Result> {
 
   protected addError(message: string, context: ParserRuleContext) {
     this.errors.push({
-      templateName: this.templateIdentifier,
+      fileIdentifier: this.templateIdentifier,
       message,
       line: context.start.line,
       column: context.start.column,

src/errors/ErrorListener.ts

@@ -5,13 +5,12 @@ import { ParserErrorItem } from "../types";
 class ErrorListener<TSymbol> extends AntlrErrorListener<TSymbol> {
   private readonly _errors: ParserErrorItem[];
 
-  constructor() {
+  constructor(public fileIdentifier: string) {
     super();
 
     this._errors = [];
   }
 
-  // TODO: improve error handling
   syntaxError(
     recognizer: Recognizer<TSymbol>,
     offendingSymbol: TSymbol,
@@ -23,8 +22,8 @@ class ErrorListener<TSymbol> extends AntlrErrorListener<TSymbol> {
       message,
       line,
       column,
+      fileIdentifier: this.fileIdentifier,
       context: null as any,
-      templateName: null,
     });
   }
 

src/index.ts

@@ -6,14 +6,20 @@ import { CircomLexer } from "./generated";
 import { ExtendedCircomParser } from "./ExtendedCircomParser";
 
 export function getCircomParser(source: string): ExtendedCircomParser {
-  const inputStream = fs.existsSync(source)
-    ? antlr4.CharStreams.fromPathSync(source, "utf8")
-    : antlr4.CharStreams.fromString(source);
+  let inputStream: antlr4.CharStream;
+  let fileIdentifier = "Built from source";
+
+  if (fs.existsSync(source)) {
+    inputStream = antlr4.CharStreams.fromPathSync(source, "utf8");
+    fileIdentifier = source;
+  } else {
+    inputStream = antlr4.CharStreams.fromString(source);
+  }
 
   const lexer = new CircomLexer(inputStream);
   const tokens = new antlr4.CommonTokenStream(lexer);
 
-  return new ExtendedCircomParser(tokens, lexer);
+  return new ExtendedCircomParser(fileIdentifier, tokens, lexer);
 }
 
 export * from "./types";

src/types/common.ts

@@ -1,3 +1,5 @@
 export type CircomValueType = bigint | CircomValueType[];
 
 export type VariableContext = Record<string, CircomValueType>;
+
+export type VariableContextWithNull = Record<string, CircomValueType | null>;

src/types/errors.ts

@@ -1,7 +1,7 @@
 import { ParserRuleContext } from "antlr4";
 
 export type ParserErrorItem = {
-  templateName: string | null;
+  fileIdentifier: string;
   message: string;
   line: number;
   column: number;

src/utils/ExpressionHelper.ts

@@ -121,10 +121,10 @@ class ExpressionVisitor extends ExtendedCircomVisitor<CircomValueType | null> {
     ctx: PIdentifierStatementContext,
   ): CircomValueType | null => {
     if (ctx.identifierStatement().idetifierAccess_list().length == 0) {
-      const variableName =
+      const variableValue =
         this.variableContext[ctx.identifierStatement().ID().getText()];
 
-      if (variableName === undefined) {
+      if (variableValue === undefined) {
         this.addError(
           `Variable ${ctx.identifierStatement().ID().getText()} is not defined`,
           ctx.identifierStatement(),
@@ -133,14 +133,36 @@ class ExpressionVisitor extends ExtendedCircomVisitor<CircomValueType | null> {
         return null;
       }
 
-      return variableName;
+      return variableValue;
     }
 
-    this.addError(
-      "IdentifierStatement is not supported with access references",
-      ctx,
-    );
-    return null;
+    const reference = ctx
+      .identifierStatement()
+      .idetifierAccess_list()
+      .map((access) => access.getText())
+      .join("");
+
+    if (reference.indexOf(".") !== -1) {
+      this.addError(
+        "IdentifierStatement is not supported with access references that are not arrays",
+        ctx,
+      );
+
+      return null;
+    }
+
+    const variableName = ctx.identifierStatement().ID().getText() + reference;
+
+    if (this.variableContext[variableName] === undefined) {
+      this.addError(
+        `Variable ${variableName} is not defined`,
+        ctx.identifierStatement(),
+      );
+
+      return null;
+    }
+
+    return this.variableContext[variableName];
   };
 
   visitPUnderscore = (_ctx: PUnderscoreContext): CircomValueType | null => {

src/utils/common.ts

@@ -1,4 +1,9 @@
 import { SimpleIdentifierListContext } from "../generated";
+import {
+  CircomValueType,
+  VariableContext,
+  VariableContextWithNull,
+} from "../types";
 
 export function parseSimpleIdentifierList(
   ctx: SimpleIdentifierListContext,
@@ -15,3 +20,115 @@ export function parseSimpleIdentifierList(
 
   return result;
 }
+
+export function buildVariableContext(
+  names: string[],
+  values: CircomValueType[],
+): VariableContext {
+  if (names.length !== values.length) {
+    throw new Error("Names and values must have the same length");
+  }
+
+  const context: VariableContext = {};
+
+  for (let i = 0; i < names.length; i++) {
+    const bindContext = bindVariableContext(
+      names[i],
+      getArrayDimensions(values[i]),
+      values[i],
+    );
+
+    for (const key in bindContext) {
+      if (bindContext[key] !== null) {
+        context[key] = bindContext[key];
+      }
+    }
+  }
+
+  return context;
+}
+
+export function getArrayDimensions(value: CircomValueType): number[] {
+  if (Array.isArray(value)) {
+    return [value.length, ...getArrayDimensions(value[0])];
+  }
+
+  return [];
+}
+
+export function bindVariableContext(
+  variableName: string,
+  dimensions: number[],
+  values: CircomValueType,
+): VariableContextWithNull {
+  const context: VariableContextWithNull = {};
+  const resolved = resolveDimensions(variableName, dimensions);
+
+  for (const variable of resolved) {
+    try {
+      context[variable] = parseVariable(
+        values,
+        variable.replace(variableName, ""),
+      );
+    } catch {
+      context[variable] = null;
+    }
+  }
+
+  return context;
+}
+
+export function resolveDimensions(
+  variableName: string,
+  dimensions: number[],
+): string[] {
+  return internalResolveDimensions([variableName], dimensions, 0);
+}
+
+// reference MUST be similar to [0][1]
+function parseVariable(
+  value: CircomValueType,
+  reference: string,
+): CircomValueType {
+  const parts = reference
+    .split("[")
+    .map((part) => part.replace("]", ""))
+    .filter((part) => part !== "")
+    .map((part) => parseInt(part));
+
+  return getReferenceValueInternal(value, parts);
+}
+
+function getReferenceValueInternal(
+  value: CircomValueType,
+  reference: number[],
+): CircomValueType {
+  if (reference.length === 0) {
+    return value;
+  }
+
+  if (!Array.isArray(value)) {
+    throw new Error("INTERNAL ERROR! Reference is invalid");
+  }
+
+  return getReferenceValueInternal(value[reference[0]], reference.slice(1));
+}
+
+function internalResolveDimensions(
+  variables: string[],
+  dimensions: number[],
+  layer: number,
+): string[] {
+  if (layer >= dimensions.length) {
+    return variables;
+  }
+
+  const result: string[] = [];
+  for (let i = 0; i < dimensions[layer]; i++) {
+    for (const variable of variables) {
+      result.push(`${variable}[${i}]`);
+    }
+  }
+
+  return internalResolveDimensions(result, dimensions, layer + 1);
+}