From 109a587654699c9e808fbc08e44027647491fe37 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Fri, 26 Apr 2019 12:03:37 -0700
Subject: [PATCH] [sledge] Fix normalization of high-degree polynomials

Summary: There was a missing case for singleton monomials of degree > 1.

Reviewed By: mbouaziz

Differential Revision: D15098810

fbshipit-source-id: e6d17d899
---
 sledge/src/llair/exp.ml      | 32 +++++++++++++++-----------------
 sledge/src/llair/exp_test.ml | 11 +++++++++++
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/sledge/src/llair/exp.ml b/sledge/src/llair/exp.ml
index 36cdd0b35..5003eb716 100644
--- a/sledge/src/llair/exp.ml
+++ b/sledge/src/llair/exp.ml
@@ -838,21 +838,15 @@ let simp_add2 typ e f = simp_add_ typ (Sum.singleton e) f
    xᵢ and the multiplicities are the exponents nᵢ. *)
 module Prod = struct
   let empty = empty_qset
-  let add exp prod = Qset.add prod exp Q.one
+
+  let add exp prod =
+    assert (match exp with Integer _ -> false | _ -> true) ;
+    Qset.add prod exp Q.one
+
   let singleton exp = add exp empty
   let union = Qset.union
 end
 
-(* map over each monomial of a polynomial *)
-let poly_map_monos poly ~f =
-  match poly with
-  | Add {typ; args= sum} ->
-      Sum.to_exp typ
-        (Sum.map sum ~f:(function
-          | Mul {typ; args= prod} -> Mul {typ; args= f prod}
-          | _ -> violates invariant poly ))
-  | _ -> fail "poly_map_monos" ()
-
 let rec simp_mul2 typ e f =
   match (e, f) with
   (* c₁ × c₂ ==> c₁×c₂ *)
@@ -869,13 +863,17 @@ let rec simp_mul2 typ e f =
   | Integer {data= c}, x | x, Integer {data= c} ->
       Sum.to_exp typ (Sum.singleton ~coeff:(Q.of_z c) x)
   (* (∏ᵤ₌₀ⁱ xᵤ) × (∏ᵥ₌ᵢ₊₁ⁿ xᵥ) ==> ∏ⱼ₌₀ⁿ xⱼ *)
-  | Mul {typ; args= xs1}, Mul {args= xs2} ->
-      Mul {typ; args= Prod.union xs1 xs2}
+  | Mul {args= xs1}, Mul {args= xs2} -> Mul {typ; args= Prod.union xs1 xs2}
   (* (∏ᵢ xᵢ) × (∑ᵤ cᵤ × ∏ⱼ yᵤⱼ) ==> ∑ᵤ cᵤ × ∏ᵢ xᵢ × ∏ⱼ yᵤⱼ *)
-  | Mul {args= prod}, (Add _ as poly) | (Add _ as poly), Mul {args= prod} ->
-      poly_map_monos ~f:(Prod.union prod) poly
+  | (Mul {args= prod} as m), Add {args= sum}
+   |Add {args= sum}, (Mul {args= prod} as m) ->
+      Sum.to_exp typ
+        (Sum.map sum ~f:(function
+          | Mul {args} -> Mul {typ; args= Prod.union prod args}
+          | Integer _ as c -> simp_mul2 typ c m
+          | mono -> Mul {typ; args= Prod.add mono prod} ))
   (* x₀ × (∏ᵢ₌₁ⁿ xᵢ) ==> ∏ᵢ₌₀ⁿ xᵢ *)
-  | Mul {typ; args= xs1}, x | x, Mul {typ; args= xs1} ->
+  | Mul {args= xs1}, x | x, Mul {args= xs1} ->
       Mul {typ; args= Prod.add x xs1}
   (* e × (∑ᵤ cᵤ × ∏ⱼ yᵤⱼ) ==> ∑ᵤ e × cᵤ × ∏ⱼ yᵤⱼ *)
   | Add {args}, e | e, Add {args} ->
@@ -899,7 +897,7 @@ let simp_negate typ x = simp_mul2 typ (minus_one typ) x
 let simp_sub typ x y =
   match (x, y) with
   (* i - j *)
-  | Integer {data= i; typ}, Integer {data= j} ->
+  | Integer {data= i}, Integer {data= j} ->
       let bits = Option.value_exn (Typ.prim_bit_size_of typ) in
       integer (Z.bsub ~bits i j) typ
   (* x - y ==> x + (-1 * y) *)
diff --git a/sledge/src/llair/exp_test.ml b/sledge/src/llair/exp_test.ml
index 05b8ba270..59c08ea18 100644
--- a/sledge/src/llair/exp_test.ml
+++ b/sledge/src/llair/exp_test.ml
@@ -254,4 +254,15 @@ let%test_module _ =
         (null = %z_2)
 
         (null = %z_2) |}]
+
+    let%expect_test _ =
+      let z1 = z + !1 in
+      let z1_2 = z1 * z1 in
+      pp z1_2 ;
+      pp (z1_2 * z1_2) ;
+      [%expect
+        {|
+        ((%z_2^2) + 2 × %z_2 + 1)
+
+        (6 × (%z_2^2) + 4 × (%z_2^3) + (%z_2^4) + 4 × %z_2 + 1) |}]
   end )