js: Add pseudo-boolean high-level functions (#7426)

* js: Add pseudo-boolean high-level functions * Add check for arg length
Z3Prover · Nov 2, 2024 · 604714b · 604714b
1 parent 91dc02d
commit 604714b
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 0 deletions.
diff --git a/src/api/js/src/high-level/high-level.test.ts b/src/api/js/src/high-level/high-level.test.ts
@@ -161,6 +161,19 @@ describe('high-level', () => {
   });
 
   describe('booleans', () => {
+    it('can use pseudo-boolean constraints', async () => {
+      const { Bool, PbEq, Solver } = api.Context('main');
+      const x = Bool.const('x');
+      const y = Bool.const('y');
+
+      const solver = new Solver();
+      solver.add(PbEq([x, y], [1, 1], 1));
+      expect(await solver.check()).toStrictEqual('sat');
+
+      solver.add(x.eq(y));
+      expect(await solver.check()).toStrictEqual('unsat');
+    });
+
     it("proves De Morgan's Law", async () => {
       const { Bool, Not, And, Eq, Or } = api.Context('main');
       const [x, y] = [Bool.const('x'), Bool.const('y')];

diff --git a/src/api/js/src/high-level/high-level.ts b/src/api/js/src/high-level/high-level.ts
@@ -961,6 +961,57 @@ export function createApi(Z3: Z3Core): Z3HighLevel {
       }
     }
 
+    function PbEq(args: [Bool<Name>, ...Bool<Name>[]], coeffs: [number, ...number[]], k: number): Bool<Name> {
+      _assertContext(...args);
+      if (args.length !== coeffs.length) {
+        throw new Error('Number of arguments and coefficients must match');
+      }
+      return new BoolImpl(
+        check(
+          Z3.mk_pbeq(
+            contextPtr,
+            args.map(arg => arg.ast),
+            coeffs,
+            k,
+          ),
+        ),
+      );
+    }
+
+    function PbGe(args: [Bool<Name>, ...Bool<Name>[]], coeffs: [number, ...number[]], k: number): Bool<Name> {
+      _assertContext(...args);
+      if (args.length !== coeffs.length) {
+        throw new Error('Number of arguments and coefficients must match');
+      }
+      return new BoolImpl(
+        check(
+          Z3.mk_pbge(
+            contextPtr,
+            args.map(arg => arg.ast),
+            coeffs,
+            k,
+          ),
+        ),
+      );
+    }
+
+    function PbLe(args: [Bool<Name>, ...Bool<Name>[]], coeffs: [number, ...number[]], k: number): Bool<Name> {
+      _assertContext(...args);
+      if (args.length !== coeffs.length) {
+        throw new Error('Number of arguments and coefficients must match');
+      }
+      return new BoolImpl(
+        check(
+          Z3.mk_pble(
+            contextPtr,
+            args.map(arg => arg.ast),
+            coeffs,
+            k,
+          ),
+        ),
+      );
+    }
+
     function ForAll<QVarSorts extends NonEmptySortArray<Name>>(
       quantifiers: ArrayIndexType<Name, QVarSorts>,
       body: Bool<Name>,
@@ -2883,6 +2934,9 @@ export function createApi(Z3: Z3Core): Z3HighLevel {
       Not,
       And,
       Or,
+      PbEq,
+      PbGe,
+      PbLe,
       ForAll,
       Exists,
       Lambda,

diff --git a/src/api/js/src/high-level/types.ts b/src/api/js/src/high-level/types.ts
@@ -429,6 +429,15 @@ export interface Context<Name extends string = 'main'> {
   /** @category Operations */
   Or(...args: Probe<Name>[]): Probe<Name>;
 
+  /** @category Operations */
+  PbEq(args: [Bool<Name>, ...Bool<Name>[]], coeffs: [number, ...number[]], k: number): Bool<Name>;
+
+  /** @category Operations */
+  PbGe(args: [Bool<Name>, ...Bool<Name>[]], coeffs: [number, ...number[]], k: number): Bool<Name>;
+
+  /** @category Operations */
+  PbLe(args: [Bool<Name>, ...Bool<Name>[]], coeffs: [number, ...number[]], k: number): Bool<Name>;
+
   // Quantifiers
 
   /** @category Operations */