Sync streams compiler: Refactor expression handling

packages/sync-rules/src/compiler/compiler.ts

@@ -6,6 +6,7 @@ import { StreamOptions, SyncPlan } from '../sync_plan/plan.js';
 import { CompilerModelToSyncPlan } from './ir_to_sync_plan.js';
 import { QuerierGraphBuilder } from './querier_graph.js';
 import { StreamQueryParser } from './parser.js';
+import { NodeLocations } from './expression.js';
 
 /**
  * State for compiling sync streams.
@@ -39,6 +40,7 @@ export class SyncStreamsCompiler {
         const parser = new StreamQueryParser({
           compiler: this,
           originalText: sql,
+          locations: new NodeLocations(),
           errors
         });
         const query = parser.parse(stmt);

packages/sync-rules/src/compiler/expression.ts

@@ -1,8 +1,12 @@
-import { Expr } from 'pgsql-ast-parser';
+import { Expr, NodeLocation, PGNode } from 'pgsql-ast-parser';
 import { SourceResultSet } from './table.js';
 import { EqualsIgnoringResultSet, equalsIgnoringResultSetList } from './compatibility.js';
 import { StableHasher } from './equality.js';
 import { ConnectionParameterSource } from '../sync_plan/plan.js';
+import { ExternalData, SqlExpression } from '../sync_plan/expression.js';
+import { ExpressionToSqlite } from '../sync_plan/expression_to_sql.js';
+import { RecursiveExpressionVisitor } from '../sync_plan/expression_visitor.js';
+import { getLocation } from '../errors.js';
 
 /**
  * An analyzed SQL expression tracking dependencies on non-static data (i.e. rows or connection sources).
@@ -16,23 +20,46 @@ import { ConnectionParameterSource } from '../sync_plan/plan.js';
  * clauses) and to evaluate expressions at runtime (by preparing them as a statement and binding external values).
  */
 export class SyncExpression implements EqualsIgnoringResultSet {
+  #sql?: string;
+  #instantiation?: readonly ExpressionInput[];
+
+  /**
+   * The original expression, where references to row or connection parameters have been replaced with SQL variables
+   * that are tracked through {@link instantiation}.
+   *
+   * This is only used to compute hash codes and to check instances for equality. {@link node} is the canonical
+   * representation of this expression.
+   */
+  get sql(): string {
+    return (this.#sql ??= ExpressionToSqlite.toSqlite(this.node));
+  }
+
+  /**
+   * The values to instantiate parameters in {@link sqlExpression} with to retain original semantics of the
+   * expression.
+   */
+  get instantiation(): readonly ExpressionInput[] {
+    if (this.#instantiation != null) {
+      return this.#instantiation;
+    }
+
+    const instantiation: ExpressionInput[] = [];
+    FindExternalData.instance.visit(this.node, instantiation);
+    return (this.#instantiation = instantiation);
+  }
+
+  get location(): NodeLocation {
+    return this.locations.locationFor(this.node);
+  }
+
   constructor(
-    /**
-     * The original expression, where references to row or connection parameters have been replaced with SQL variables
-     * that are tracked through {@link instantiation}.
-     */
-    readonly sql: string,
     /**
      * The AST node backing {@link sql}.
      *
      * We use this to be able to compose expressions, e.g. to possibly merge them.
      */
-    readonly node: Expr,
-    /**
-     * The values to instantiate parameters in {@link sqlExpression} with to retain original semantics of the
-     * expression.
-     */
-    readonly instantiation: ExpressionInputWithSpan[]
+    readonly node: SqlExpression<ExpressionInput>,
+    readonly locations: NodeLocations
   ) {}
 
   equalsAssumingSameResultSet(other: EqualsIgnoringResultSet): boolean {
@@ -47,32 +74,22 @@ export class SyncExpression implements EqualsIgnoringResultSet {
     hasher.addString(this.sql);
     equalsIgnoringResultSetList.hash(hasher, this.instantiation);
   }
-
-  *instantiationValues() {
-    for (const instantiation of this.instantiation) {
-      yield instantiation.value;
-    }
-  }
 }
 
-export type ExpressionInput = ColumnInRow | ConnectionParameter;
-
-export class ExpressionInputWithSpan implements EqualsIgnoringResultSet {
-  constructor(
-    readonly value: ExpressionInput,
-    readonly startOffset: number,
-    readonly length: number
-  ) {}
-
-  equalsAssumingSameResultSet(other: EqualsIgnoringResultSet): boolean {
-    return other instanceof ExpressionInputWithSpan && other.value.equalsAssumingSameResultSet(this.value);
+class FindExternalData extends RecursiveExpressionVisitor<ExpressionInput, void, ExpressionInput[]> {
+  defaultExpression(expr: SqlExpression<ExpressionInput>, arg: ExpressionInput[]): void {
+    this.visitChildren(expr, arg);
   }
 
-  assumingSameResultSetEqualityHashCode(hasher: StableHasher): void {
-    return this.value.assumingSameResultSetEqualityHashCode(hasher);
+  visitExternalData(expr: ExternalData<ExpressionInput>, arg: ExpressionInput[]): void {
+    arg.push(expr.source);
   }
+
+  static readonly instance: FindExternalData = new FindExternalData();
 }
 
+export type ExpressionInput = ColumnInRow | ConnectionParameter;
+
 export class ColumnInRow implements EqualsIgnoringResultSet {
   constructor(
     readonly syntacticOrigin: Expr,
@@ -103,3 +120,23 @@ export class ConnectionParameter implements EqualsIgnoringResultSet {
     hasher.addString(this.source);
   }
 }
+
+/**
+ * Tracks the original source location for translated {@link SqlExpression} nodes.
+ *
+ * We want to serialize translated expressions for sync plan, so embedding source offsets in them expands the size of
+ * sync plans and is tedious. We only need access to node locations while compiling sync streams, which we store in this
+ * in-memory map.
+ */
+export class NodeLocations {
+  readonly sourceForNode = new Map<SqlExpression<ExpressionInput>, PGNode | NodeLocation>();
+
+  locationFor(source: SqlExpression<ExpressionInput>): NodeLocation {
+    const location = getLocation(this.sourceForNode.get(source));
+    if (location == null) {
+      throw new Error('Missing location');
+    }
+
+    return location;
+  }
+}

packages/sync-rules/src/compiler/filter.ts

@@ -45,7 +45,7 @@ export class SingleDependencyExpression implements EqualsIgnoringResultSet {
 
   constructor(expression: SyncExpression | SingleDependencyExpression) {
     if (expression instanceof SyncExpression) {
-      const checked = SingleDependencyExpression.extractSingleDependency(expression.instantiationValues());
+      const checked = SingleDependencyExpression.extractSingleDependency(expression.instantiation);
       if (checked == null) {
         throw new InvalidExpressionError('Expression with multiple dependencies passed to SingleDependencyExpression');
       }
@@ -133,7 +133,7 @@ export class EqualsClause {
   ) {}
 
   get location(): NodeLocation | undefined {
-    return expandNodeLocations([this.left.expression.node, this.right.expression.node]);
+    return expandNodeLocations([this.left.expression.location, this.right.expression.location]);
   }
 }
 

packages/sync-rules/src/compiler/filter_simplifier.ts

@@ -1,9 +1,8 @@
-import { assignChanged, astMapper, BinaryOperator, Expr } from 'pgsql-ast-parser';
 import { And, BaseTerm, EqualsClause, isBaseTerm, Or, SingleDependencyExpression } from './filter.js';
-import { PostgresToSqlite } from './sqlite.js';
-import { ExpressionInput, SyncExpression } from './expression.js';
-import { expandNodeLocations } from '../errors.js';
+import { SyncExpression } from './expression.js';
 import { SourceResultSet } from './table.js';
+import { BinaryOperator } from '../sync_plan/expression.js';
+import { expandNodeLocations } from '../errors.js';
 
 export class FilterConditionSimplifier {
   constructor(private readonly originalText: string) {}
@@ -21,7 +20,7 @@ export class FilterConditionSimplifier {
       }
     }
 
-    baseTerms = this.mergeByCommonDependencies('OR', baseTerms);
+    baseTerms = this.mergeByCommonDependencies('or', baseTerms);
     for (const term of baseTerms) {
       andTerms.push({ terms: [term] });
     }
@@ -30,7 +29,7 @@ export class FilterConditionSimplifier {
   }
 
   private simplifyAnd(and: And): And | BaseTerm {
-    const merged = this.mergeByCommonDependencies('AND', and.terms);
+    const merged = this.mergeByCommonDependencies('and', and.terms);
 
     if (merged.length == 1) {
       return merged[0];
@@ -92,8 +91,8 @@ export class FilterConditionSimplifier {
       // must be a row condition since it can't be represented as parameters that could be instantiated.
       if (
         SingleDependencyExpression.extractSingleDependency([
-          ...base.left.expression.instantiationValues(),
-          ...base.right.expression.instantiationValues()
+          ...base.left.expression.instantiation,
+          ...base.right.expression.instantiation
         ])
       ) {
         return this.composeExpressions('=', base.left, base.right);
@@ -109,52 +108,26 @@ export class FilterConditionSimplifier {
    * For instance, `composeExpressions('AND', a, b, c)` returns `a AND b AND c` as a single expression. All expressions
    * must have compatible dependencies.
    */
-  private composeExpressions(op: BinaryOperator, ...terms: SingleDependencyExpression[]): SingleDependencyExpression {
+  private composeExpressions(
+    operator: BinaryOperator,
+    ...terms: SingleDependencyExpression[]
+  ): SingleDependencyExpression {
     if (terms.length == 0) {
       throw new Error("Can't compose zero expressions");
     }
 
-    let node: Expr | null = null;
-    const instantiation: ExpressionInput[] = [];
-    const transformer = astMapper(() => ({
-      parameter: (st) => {
-        // All parameters are named ?<idx>, increase the index to avoid collisions with parameters we've already added.
-        const originalIndex = Number(st.name.substring(1));
-        const newIndex = instantiation.length + originalIndex;
-
-        return assignChanged(st, { name: `?${newIndex}` });
-      }
-    }));
-
-    for (const element of terms) {
-      if (node == null) {
-        node = element.expression.node;
-      } else {
-        const transformed = transformer.expr(element.expression.node)!;
-
-        node = {
-          type: 'binary',
-          op,
-          left: node,
-          right: transformed,
-          _location: expandNodeLocations([node, transformed])
-        };
-      }
+    const [first, ...rest] = terms;
+    const locations = first.expression.locations;
+    let inner = first.expression.node;
+    for (const additional of rest) {
+      inner = { type: 'binary', operator, left: inner, right: additional.expression.node };
+    }
 
-      instantiation.push(...element.expression.instantiationValues());
+    const location = expandNodeLocations(terms.map((e) => e.expression.location));
+    if (location) {
+      locations.sourceForNode.set(inner, location);
     }
 
-    const toSqlite = new PostgresToSqlite(
-      this.originalText,
-      {
-        report() {
-          // We don't need to re-report errors when we shuffle expressions around, the mapper would have already reported
-          // these issues on the first round.
-        }
-      },
-      instantiation
-    );
-    toSqlite.addExpression(node!);
-    return new SingleDependencyExpression(new SyncExpression(toSqlite.sql, node!, toSqlite.inputs));
+    return new SingleDependencyExpression(new SyncExpression(inner, locations));
   }
 }

packages/sync-rules/src/compiler/ir_to_sync_plan.ts

@@ -1,9 +1,11 @@
 import * as plan from '../sync_plan/plan.js';
+import { SqlExpression } from '../sync_plan/expression.js';
 import * as resolver from './bucket_resolver.js';
 import { CompiledStreamQueries } from './compiler.js';
 import { Equality, HashMap, StableHasher, unorderedEquality } from './equality.js';
-import { ColumnInRow, SyncExpression } from './expression.js';
+import { ColumnInRow, ExpressionInput, SyncExpression } from './expression.js';
 import * as rows from './rows.js';
+import { MapSourceVisitor, visitExpr } from '../sync_plan/expression_visitor.js';
 
 export class CompilerModelToSyncPlan {
   private static readonly evaluatorHash: Equality<rows.RowEvaluator[]> = unorderedEquality({
@@ -121,19 +123,16 @@ export class CompilerModelToSyncPlan {
     });
   }
 
-  private translateExpression<T extends plan.SqlParameterValue>(expression: SyncExpression): plan.SqlExpression<T> {
-    return {
-      sql: expression.sql,
-      values: expression.instantiation.map((e) => {
-        const value = e.value;
-
-        if (value instanceof ColumnInRow) {
-          return { column: value.column } satisfies plan.ColumnSqlParameterValue;
-        } else {
-          return { request: value.source } satisfies plan.RequestSqlParameterValue;
-        }
-      }) as unknown[] as T[]
-    };
+  private translateExpression<T extends plan.SqlParameterValue>(expression: SyncExpression): SqlExpression<T> {
+    const mapper = new MapSourceVisitor<ExpressionInput, T>((value) => {
+      if (value instanceof ColumnInRow) {
+        return { column: value.column } satisfies plan.ColumnSqlParameterValue as unknown as T;
+      } else {
+        return { request: value.source } satisfies plan.RequestSqlParameterValue as unknown as T;
+      }
+    });
+
+    return visitExpr(mapper, expression.node, null);
   }
 
   private translateStreamResolver(value: resolver.StreamResolver): plan.StreamQuerier {