Building a CAS in OCaml Part 3

expression tree            polynomial
-----------------------    -----------------------------------
sin(x) + 1/x + (a+b)^7     a sorted list of (exponents, coefficient)
arbitrary nesting          no nesting at all
simplify is a rewrite      add is a merge, mul is a double loop
equality needs a canon-    equality is list equality
  ical form                
 
expand, gcd, resultant, coeff, degree, diff
  all live on the right and are exported on the left

type monom = int array
type term = monom * Rational.t
 
type t = {
  vars : string array; (* sorted, no duplicates *)
  terms : term list;   (* descending lex, no zero coefficients *)
}

type order = Lex | Grlex | Grevlex
 
let total_deg (m : monom) : int = Array.fold_left ( + ) 0 m
 
let cmp_lex (a : monom) (b : monom) : int =
  let n = Array.length a in
  let rec go i = if i >= n then 0 else if a.(i) <> b.(i) then compare a.(i) b.(i) else go (i + 1) in
  go 0
 
let cmp_grlex (a : monom) (b : monom) : int =
  let c = compare (total_deg a) (total_deg b) in
  if c <> 0 then c else cmp_lex a b
 
(* Graded reverse lex: equal degree is broken by the LAST variable in
   which the exponents differ, with the smaller exponent winning. It is
   the order Groebner basis computations are fastest in, which is why
   it is here before anything needs it. *)
let cmp_grevlex (a : monom) (b : monom) : int =
  let c = compare (total_deg a) (total_deg b) in
  if c <> 0 then c
  else begin
    let n = Array.length a in
    let rec go i =
      if i < 0 then 0 else if a.(i) <> b.(i) then compare b.(i) a.(i) else go (i - 1)
    in
    go (n - 1)
  end
 
let cmp_order (o : order) : monom -> monom -> int =
  match o with Lex -> cmp_lex | Grlex -> cmp_grlex | Grevlex -> cmp_grevlex

let mono_mul (a : monom) (b : monom) : monom = Array.init (Array.length a) (fun i -> a.(i) + b.(i))
 
(* Some (a / b) when b divides a, that is when no exponent goes
   negative. This single test is the whole of multivariate division. *)
let mono_div (a : monom) (b : monom) : monom option =
  let n = Array.length a in
  let r = Array.make n 0 in
  let ok = ref true in
  for i = 0 to n - 1 do
    let d = a.(i) - b.(i) in
    if d < 0 then ok := false else r.(i) <- d
  done;
  if !ok then Some r else None

  (* The textbook case that separates grlex from grevlex: same degree,
     and they disagree. *)
  let a = [| 1; 2; 0 |] and b = [| 0; 0; 3 |] in
  check "grlex prefers the earlier variable" (cmp_grlex a b > 0);
  check "grevlex prefers the later variable being absent" (cmp_grevlex a b > 0);
  let a = [| 1; 1; 1 |] and b = [| 2; 0; 1 |] in
  check "grlex: x^2*z beats x*y*z" (cmp_grlex a b < 0);
  check "grevlex: x^2*z beats x*y*z too" (cmp_grevlex a b < 0)

(* Sort descending, add up equal monomials, and drop the zeros. *)
let make (vars : string array) (terms : term list) : t =
  let sorted = List.sort (fun (m1, _) (m2, _) -> cmp_lex m2 m1) terms in
  let rec collect = function
    | (m1, c1) :: (m2, c2) :: rest when cmp_lex m1 m2 = 0 ->
      collect ((m1, Rational.add c1 c2) :: rest)
    | x :: rest -> x :: collect rest
    | [] -> []
  in
  { vars; terms = List.filter (fun (_, c) -> not (Rational.is_zero c)) (collect sorted) }
 
let zero : t = { vars = [||]; terms = [] }
 
let of_rational (r : Rational.t) : t =
  if Rational.is_zero r then zero else { vars = [||]; terms = [ ([||], r) ] }
 
let of_int (n : int) : t = of_rational (Rational.of_int n)
let one : t = of_int 1
let var (name : string) : t = { vars = [| name |]; terms = [ ([| 1 |], Rational.one) ] }

(* Re-express p over a superset of its own variables. The index map is
   computed once and then applied to every monomial. *)
let remap (p : t) (vars : string array) : t =
  if p.vars = vars then p
  else begin
    let n = Array.length vars in
    let idx =
      Array.map
        (fun v ->
          let rec find i = if i >= n then -1 else if vars.(i) = v then i else find (i + 1) in
          find 0)
        p.vars
    in
    Array.iter (fun i -> if i < 0 then invalid_arg "Poly.remap: missing variable") idx;
    let terms =
      List.map
        (fun (m, c) ->
          let m' = Array.make n 0 in
          Array.iteri (fun j e -> m'.(idx.(j)) <- e) m;
          (m', c))
        p.terms
    in
    make vars terms
  end
 
let union_vars (a : string array) (b : string array) : string array =
  let l = List.sort_uniq String.compare (Array.to_list a @ Array.to_list b) in
  Array.of_list l
 
let unify (a : t) (b : t) : t * t =
  if a.vars = b.vars then (a, b)
  else begin
    let v = union_vars a.vars b.vars in
    (remap a v, remap b v)
  end

(* Note there is no pruning step: a polynomial's variable array is
   allowed to be a superset of the variables it actually uses, which is
   what happens whenever a term cancels. Keeping it that way means
   every monomial in a polynomial is indexed the same, so the division
   loop below can hand monomials between two polynomials without
   re-aligning them mid-loop. `degree_in` reports 0 for a variable the
   terms never mention, which is the right answer anyway. *)

let neg (p : t) : t = { p with terms = List.map (fun (m, c) -> (m, Rational.neg c)) p.terms }
 
let add (a : t) (b : t) : t =
  let a, b = unify a b in
  make a.vars (a.terms @ b.terms)
 
let sub (a : t) (b : t) : t = add a (neg b)
 
let scale (p : t) (c : Rational.t) : t =
  if Rational.is_zero c then zero
  else { p with terms = List.map (fun (m, x) -> (m, Rational.mul c x)) p.terms }
 
let mul (a : t) (b : t) : t =
  if is_zero a || is_zero b then zero
  else begin
    let a, b = unify a b in
    let acc = ref [] in
    List.iter
      (fun (m1, c1) ->
        List.iter (fun (m2, c2) -> acc := (mono_mul m1 m2, Rational.mul c1 c2) :: !acc) b.terms)
      a.terms;
    make a.vars !acc
  end

let pow (p : t) (n : int) : t =
  if n < 0 then invalid_arg "Poly.pow: negative exponent";
  let rec go acc b n =
    if n = 0 then acc else if n land 1 = 1 then go (mul acc b) (mul b b) (n lsr 1) else go acc (mul b b) (n lsr 1)
  in
  go one p n

(* Degree in one variable; -1 for the zero polynomial, 0 for anything
   that does not mention the variable at all. *)
let degree_in (p : t) (v : string) : int =
  if is_zero p then -1
  else
    match var_index p v with
    | None -> 0
    | Some i -> List.fold_left (fun d (m, _) -> max d m.(i)) 0 p.terms
 
let total_degree (p : t) : int =
  if is_zero p then -1 else List.fold_left (fun d (m, _) -> max d (total_deg m)) 0 p.terms
 
(* The coefficient of v^k, itself a polynomial in the other variables. *)
let coeff_in (p : t) (v : string) (k : int) : t =
  if is_zero p then zero
  else
    match var_index p v with
    | None -> if k = 0 then p else zero
    | Some i ->
      let terms =
        List.filter_map
          (fun (m, c) ->
            if m.(i) <> k then None
            else begin
              let m' = Array.copy m in
              m'.(i) <- 0;
              Some (m', c)
            end)
          p.terms
      in
      make p.vars terms
 
(* Coefficients by ascending degree in v: index k holds the coefficient
   of v^k. This is the view every univariate algorithm below works in. *)
let coeffs_in (p : t) (v : string) : t list =
  let d = degree_in p v in
  if d < 0 then [] else List.init (d + 1) (fun k -> coeff_in p v k)
 
let of_coeffs (v : string) (cs : t list) : t =
  let x = var v in
  let rec go k acc = function
    | [] -> acc
    | c :: rest -> go (k + 1) (add acc (mul c (pow x k))) rest
  in
  go 0 zero cs
 
(* Leading coefficient with respect to v, as a polynomial. *)
let lc_in (p : t) (v : string) : t =
  let d = degree_in p v in
  if d < 0 then zero else coeff_in p v d

degree_in ((x + y)^3) x        =  3
coeffs_in ((x + y)^3) x        =  y^3 | 3*y^2 | 3*y | 1
of_coeffs x (coeffs_in p x)    =  p
degree_in zero x               = -1
degree_in (y^5) x              =  0

exception Not_exact
 
(* Exact division. Because the coefficients form a field, the only way
   a step can fail is a leading monomial that does not divide, and if
   b really divides a that never happens - the quotient's leading term
   has to be lt(a)/lt(b). So the first failure is a proof that the
   division is not exact, and there is no need to accumulate a
   remainder. *)
let divide (a : t) (b : t) : t option =
  if is_zero b then raise Division_by_zero
  else if is_zero a then Some zero
  else begin
    let a, b = unify a b in
    let bm, bc = List.hd b.terms in
    let quo = ref [] and rem = ref a in
    try
      while not (is_zero !rem) do
        let rm, rc = List.hd !rem.terms in
        match mono_div rm bm with
        | None -> raise Not_exact
        | Some m ->
          let c = Rational.div rc bc in
          quo := (m, c) :: !quo;
          rem := sub !rem (mul { vars = a.vars; terms = [ (m, c) ] } b)
      done;
      Some (make a.vars !quo)
    with Not_exact -> None
  end

let divide_exn (a : t) (b : t) : t =
  match divide a b with Some q -> q | None -> failwith "Poly.divide_exn: not an exact division"
 
let divides (b : t) (a : t) : bool = match divide a b with Some _ -> true | None -> false

(* Pseudo-division with respect to v, treating both arguments as
   univariate in v over the ring of polynomials in the other
   variables. Exactly the Z[x] story from Upoly, one level up. *)
let pseudo_div (v : string) (a : t) (b : t) : t * t =
  let db = degree_in b v in
  if is_zero b then raise Division_by_zero;
  let da = degree_in a v in
  if da < db then (zero, a)
  else begin
    let lcb = lc_in b v in
    let x = var v in
    let e = ref (da - db + 1) in
    let q = ref zero and r = ref a in
    while (not (is_zero !r)) && degree_in !r v >= db do
      let d = degree_in !r v - db in
      let c = lc_in !r v in
      let term = mul c (pow x d) in
      q := add (mul !q lcb) term;
      r := sub (mul !r lcb) (mul term b);
      decr e
    done;
    let f = pow lcb !e in
    (mul !q f, mul !r f)
  end

  (* Pseudo-division identity, one variable at a time. *)
  let a = add (mul (pow x 3) y) (add (mul x y) (n 5)) in
  let b = add (mul (pow x 2) y) (n 1) in
  let q, r = pseudo_div "x" a b in
  let e = degree_in a "x" - degree_in b "x" + 1 in
  check "pseudo division identity"
    (equal (add (mul q b) r) (mul (pow (lc_in b "x") e) a));
  check "pseudo remainder degree drops" (degree_in r "x" < degree_in b "x")

./_out/main growth

a = x^8 + x^6 - 3*x^4 - 3*x^3 + 8*x^2 + 2*x - 5
b = 3*x^6 + 5*x^4 - 4*x^2 - 9*x + 21
  prem 1: degree  6, max coefficient 21
  prem 2: degree  4, max coefficient 15
  prem 3: degree  2, max coefficient 59535
  prem 4: degree  1, max coefficient 1654608338437500
  prem 5: degree  0, max coefficient 12593338795500743100931141992187500
subresultant gcd = 1

(* Knuth's Algorithm C. The plain pseudo-remainder sequence is correct
   but its coefficients grow doubly exponentially; the g and h below
   track a factor that provably divides the next remainder exactly, and
   dividing it out holds the growth to something linear in the degree.
   `steps` counts the remainders taken, for the benchmark. *)
let steps = ref 0
 
let gcd_prs (a : t) (b : t) : t =
  if is_zero a then primitive_part b
  else if is_zero b then primitive_part a
  else begin
    let d = Bigint.gcd (content a) (content b) in
    let u = ref (primitive_part a) and v = ref (primitive_part b) in
    if degree !u < degree !v then begin
      let t = !u in
      u := !v;
      v := t
    end;
    let g = ref Bigint.one and h = ref Bigint.one in
    let result = ref zero and finished = ref false in
    while not !finished do
      incr steps;
      let delta = degree !u - degree !v in
      let r = prem !u !v in
      if is_zero r then begin
        result := !v;
        finished := true
      end
      else if degree r = 0 then begin
        result := one;
        finished := true
      end
      else begin
        u := !v;
        v := divexact_int r (Bigint.mul !g (Bigint.pow !h delta));
        g := lc !u;
        h :=
          if delta = 0 then !h
          else Bigint.div (Bigint.pow !g delta) (Bigint.pow !h (delta - 1))
      end
    done;
    scale (primitive_part !result) d
  end

(* Pseudo-division: returns (q, r) with lc(v)^(deg u - deg v + 1) * u =
   q*v + r and deg r < deg v. Multiplying through by that power of the
   leading coefficient is what keeps every intermediate value in Z
   without ever needing a fraction. *)
let pseudo_div (u : t) (v : t) : t * t =
  if is_zero v then raise Division_by_zero;
  let dv = degree v and lcv = lc v in
  if degree u < dv then (zero, u)
  else begin
    let e = ref (degree u - dv + 1) in
    let q = ref zero and r = ref u in
    while (not (is_zero !r)) && degree !r >= dv do
      let d = degree !r - dv in
      let c = lc !r in
      let term = shift (const c) d in
      q := add (scale !q lcv) term;
      r := sub (scale !r lcv) (mul term v);
      decr e
    done;
    (* The loop consumed one factor of lc(v) per step; top up so that
       the identity holds with exactly deg u - deg v + 1 factors. *)
    let f = Bigint.pow lcv !e in
    (scale !q f, scale !r f)
  end

gcd_prs (x^2 - 1) (x^2 + 2x + 1)       =  x + 1
gcd_prs (x^4 - 1) (x^3 - 1)            =  x - 1
gcd_prs (6x^2 + 18x + 12) (4x^2 + 12x + 8)
                                       =  2x^2 + 6x + 4
 
the integer content is part of the answer: gcd(6, 4) = 2

module Zp = struct
  let add p a b =
    let s = a + b in
    if s >= p then s - p else s
 
  let sub p a b =
    let s = a - b in
    if s < 0 then s + p else s
 
  let mul p a b = a * b mod p
  let neg p a = if a = 0 then 0 else p - a
 
  (* Extended Euclid, iterative. *)
  let inv p a =
    if a = 0 then invalid_arg "Zp.inv: zero";
    let r0 = ref p and r1 = ref (a mod p) in
    let t0 = ref 0 and t1 = ref 1 in
    while !r1 <> 0 do
      let q = !r0 / !r1 in
      let r = !r0 - (q * !r1) in
      r0 := !r1;
      r1 := r;
      let t = !t0 - (q * !t1) in
      t0 := !t1;
      t1 := t
    done;
    let t = !t0 mod p in
    if t < 0 then t + p else t
end

module Pmod = struct
  type t = int array
 
  let normalize (a : t) : t =
    let n = ref (Array.length a) in
    while !n > 0 && a.(!n - 1) = 0 do
      decr n
    done;
    if !n = Array.length a then a else Array.sub a 0 !n
 
  let degree (a : t) : int = Array.length a - 1
  let is_zero (a : t) : bool = Array.length a = 0
  let lc (a : t) : int = if is_zero a then 0 else a.(Array.length a - 1)
 
  let scale p (a : t) (c : int) : t =
    if c = 0 then [||] else Array.map (fun x -> Zp.mul p x c) a
 
  let monic p (a : t) : t = if is_zero a then a else scale p a (Zp.inv p (lc a))
 
  let sub p (a : t) (b : t) : t =
    let n = max (Array.length a) (Array.length b) in
    normalize
      (Array.init n (fun i ->
         let x = if i < Array.length a then a.(i) else 0 in
         let y = if i < Array.length b then b.(i) else 0 in
         Zp.sub p x y))
 
  (* Remainder of a divided by b, both over the field F_p, so no
     pseudo-division is needed. *)
  let rem p (a : t) (b : t) : t =
    if is_zero b then raise Division_by_zero;
    let db = degree b and ib = Zp.inv p (lc b) in
    let r = ref a in
    while (not (is_zero !r)) && degree !r >= db do
      let d = degree !r - db in
      let c = Zp.mul p (lc !r) ib in
      let shifted = Array.make (Array.length b + d) 0 in
      Array.iteri (fun i x -> shifted.(i + d) <- Zp.mul p x c) b;
      r := sub p !r shifted
    done;
    !r
 
  let gcd p (a : t) (b : t) : t =
    let u = ref a and v = ref b in
    while not (is_zero !v) do
      let r = rem p !u !v in
      u := !v;
      v := r
    done;
    monic p !u
 
  let of_upoly p (a : Bigint.t array) : t =
    normalize (Array.map (fun c -> bigint_mod_int c p) a)
end

(* Every coefficient of a divisor of a is bounded by
   2^deg(a) * ||a||_inf, so a GCD scaled to have leading coefficient
   gcd(lc a, lc b) is bounded by that times the scaling. Once the
   product of the primes used exceeds twice the bound, the symmetric
   CRT lift is the answer and not merely congruent to it. *)
let landau_mignotte (a : t) (b : t) (g : Bigint.t) : Bigint.t =
  let m = min (degree a) (degree b) in
  let na = max_norm a and nb = max_norm b in
  let smaller = if Bigint.compare na nb <= 0 then na else nb in
  Bigint.mul (Bigint.mul (Bigint.pow (Bigint.of_int 2) m) g) smaller
 
(* Symmetric representative of r modulo m: the one in (-m/2, m/2]. *)
let symmetric (r : Bigint.t) (m : Bigint.t) : Bigint.t =
  let two_r = Bigint.mul (Bigint.of_int 2) r in
  if Bigint.compare two_r m > 0 then Bigint.sub r m else r
 
(* One coefficient of the Chinese remainder step: find the value
   congruent to r1 mod m1 and to r2 mod p. *)
let crt_coeff (r1 : Bigint.t) (m1 : Bigint.t) (r2 : int) (p : int) : Bigint.t =
  let r1p = bigint_mod_int r1 p in
  let m1p = bigint_mod_int m1 p in
  let delta = Zp.mul p (Zp.sub p r2 r1p) (Zp.inv p m1p) in
  Bigint.add r1 (Bigint.mul m1 (Bigint.of_int delta))

let gcd_modular (a : t) (b : t) : t =
  if is_zero a then primitive_part b
  else if is_zero b then primitive_part a
  else begin
    let d = Bigint.gcd (content a) (content b) in
    let aa = primitive_part a and bb = primitive_part b in
    let g = Bigint.gcd (lc aa) (lc bb) in
    let bound = Bigint.mul (Bigint.of_int 2) (landau_mignotte aa bb g) in
    (* v holds the CRT lift so far, m the product of the primes used. *)
    let v = ref [||] and m = ref Bigint.one and have = ref false in
    let p = ref (1 lsl 29) in
    let result = ref zero and finished = ref false in
    while not !finished do
      p := next_prime (!p + 1);
      let pi = !p in
      (* Skip a prime that divides the leading coefficients: modulo
         such a prime the degree drops and the image is not the image
         of the GCD. *)
      if bigint_mod_int g pi <> 0 then begin
        let ap = Pmod.of_upoly pi aa and bp = Pmod.of_upoly pi bb in
        if Pmod.degree ap = degree aa && Pmod.degree bp = degree bb then begin
          incr primes_used;
          let gp = Pmod.gcd pi ap bp in
          if Pmod.degree gp = 0 then begin
            (* Coprime mod one prime means coprime over Z. *)
            result := scale one d;
            finished := true
          end
          else begin
            (* Impose the known leading coefficient instead of leaving
               the image monic, so the images agree across primes. *)
            let gp = Pmod.scale pi gp (bigint_mod_int g pi) in
            let cur_deg = if !have then Array.length !v - 1 else max_int in
            if (not !have) || Pmod.degree gp < cur_deg then begin
              (* A smaller degree means every earlier prime was
                 unlucky; throw the accumulation away and restart. *)
              have := true;
              m := Bigint.of_int pi;
              v :=
                Array.map
                  (fun c -> symmetric (Bigint.of_int c) (Bigint.of_int pi))
                  gp
            end
            else if Pmod.degree gp = cur_deg then begin
              let m' = Bigint.mul !m (Bigint.of_int pi) in
              let lifted =
                Array.init (Array.length !v) (fun i ->
                  symmetric (crt_coeff !v.(i) !m gp.(i) pi) m')
              in
              v := lifted;
              m := m'
            end;
            (* else this prime is unlucky: its GCD is too big, drop it *)
            if Bigint.compare !m bound > 0 then begin
              let cand = primitive_part (normalize (Array.copy !v)) in
              match (divides cand aa, divides cand bb) with
              | Some _, Some _ ->
                result := scale cand d;
                finished := true
              | _ -> () (* bound was not enough yet; keep going *)
            end
          end
        end
      end
    done;
    !result
  end

./_out/main bench

deg(g)   deg(a)        prs (s)  modular (s) max digits
4        8              0.0004       0.0006         19
8        16             0.0042       0.0006         18
16       32             0.0539       0.0008         19
24       48             0.2697       0.0011         19
32       64             0.7688       0.0016         19
40       80             1.7734       0.0023         19
48       96             3.5756       0.0029         19

let rec gcd (a : t) (b : t) : t =
  if is_zero a then primitive b
  else if is_zero b then primitive a
  else begin
    let a, b = unify (primitive a) (primitive b) in
    (* Recurse on the first variable either side actually uses. A
       variable handled here never reappears further down, so the
       recursion is bounded by the number of variables. *)
    let n = Array.length a.vars in
    let rec pick i =
      if i >= n then None
      else if degree_in a a.vars.(i) > 0 || degree_in b a.vars.(i) > 0 then Some a.vars.(i)
      else pick (i + 1)
    in
    match pick 0 with
    | None -> one (* both are nonzero constants, and over Q every one of those is a unit *)
    | Some v ->
      let ca = content_in v a and cb = content_in v b in
      let cg = gcd ca cb in
      let pa = divide_exn a ca and pb = divide_exn b cb in
      if degree_in pa v = 0 || degree_in pb v = 0 then primitive cg
      else primitive (mul cg (gcd_prs v pa pb))
  end
 
and content_in (v : string) (p : t) : t =
  List.fold_left (fun g c -> gcd g c) zero (coeffs_in p v)
 
and primitive_part_in (v : string) (p : t) : t = divide_exn p (content_in v p)
 
(* Knuth's Algorithm C again, this time over the ring of polynomials in
   the variables other than v. Both arguments must already be primitive
   with respect to v. *)
and gcd_prs (v : string) (a : t) (b : t) : t =
  let u = ref a and w = ref b in
  if degree_in !u v < degree_in !w v then begin
    let t = !u in
    u := !w;
    w := t
  end;
  let g = ref one and h = ref one in
  let result = ref one and finished = ref false in
  while not !finished do
    let delta = degree_in !u v - degree_in !w v in
    let r = pseudo_rem v !u !w in
    if is_zero r then begin
      result := primitive_part_in v !w;
      finished := true
    end
    else if degree_in r v = 0 then begin
      result := one;
      finished := true
    end
    else begin
      u := !w;
      w := divide_exn r (mul !g (pow !h delta));
      g := lc_in !u v;
      h := (if delta = 0 then !h else divide_exn (pow !g delta) (pow !h (delta - 1)))
    end
  done;
  !result

(* Scale by the unique positive rational that clears all denominators,
   divides out the integer content, and makes the lex-leading
   coefficient positive. A GCD is only defined up to a unit, and over Q
   every nonzero rational is a unit, so a convention like this is what
   makes `gcd` a function rather than a relation. *)
let primitive (p : t) : t =
  if is_zero p then p
  else begin
    let den_lcm =
      List.fold_left
        (fun l (_, c) ->
          let d = Rational.den c in
          Bigint.div (Bigint.mul l d) (Bigint.gcd l d))
        Bigint.one p.terms
    in
    let scaled = scale p (Rational.of_bigint den_lcm) in
    let num_gcd =
      List.fold_left (fun g (_, c) -> Bigint.gcd g (Rational.num c)) Bigint.zero scaled.terms
    in
    let r = scale scaled (Rational.make Bigint.one num_gcd) in
    match lead_term r with
    | Some (_, c) when Rational.sign c < 0 -> neg r
    | _ -> r
  end

> gcd x^6 - 1, x^4 - 1
-1 + x^2
> gcd x^2*y - y^3, x^2 - 2*x*y + y^2
x - y
> gcd x + 1, x + 2
1

let rec resultant (v : string) (a : t) (b : t) : t =
  if is_zero a || is_zero b then zero
  else begin
    let m = degree_in a v and n = degree_in b v in
    if m = 0 && n = 0 then one
    else if m < n then
      let s = resultant v b a in
      if m * n mod 2 = 1 then neg s else s
    else if n = 0 then pow b m
    else begin
      let r = pseudo_rem v a b in
      if is_zero r then zero
      else begin
        let k = degree_in r v in
        let lcb = lc_in b v in
        let num = mul (resultant v b r) (pow lcb (m - k)) in
        let den = pow lcb (n * (m - n + 1)) in
        let s = divide_exn num den in
        if m * n mod 2 = 1 then neg s else s
      end
    end
  end

(* res(a, b) = 0 exactly when a and b have a common factor. *)
let test_resultant () =
  check_str "resultant eliminates x" "y^2 - 7*y + 9"
    (s (resultant "x" (sub (pow x 2) y) (add (sub (pow x 2) (n 3)) x)));
  check_str "resultant with a linear argument" "-y + z^2"
    (s (resultant "x" (sub (pow x 2) y) (sub x z)));
  check "a shared root makes the resultant vanish"
    (is_zero (resultant "x" (mul (add x (n 1)) (add x (n 2))) (mul (add x (n 1)) (add x (n 3)))));
  check "no shared root makes it nonzero"
    (not (is_zero (resultant "x" (add x (n 1)) (add x (n 2)))));
  (* The discriminant of a x^2 + b x + c is b^2 - 4ac, and it is
     res(p, p') up to the leading coefficient. *)
  (* For p = z*x^2 + y*x + c the discriminant is y^2 - 4*z*c, and
     res(p, p') = -z * (y^2 - 4*z*c), so dividing by -4 gives
     (1/4)*y^2*z - c*z^2. *)
  let p = add (add (mul z (pow x 2)) (mul y x)) (var "c") in
  let r = resultant "x" p (diff p "x") in
  check_str "discriminant of a quadratic" "-c*z^2 + 1/4*y^2*z"
    (s (scale r (Rational.make (Bigint.of_int (-1)) (Bigint.of_int 4))));
  (* res(a, b) = (-1)^(deg a * deg b) res(b, a): both degrees are 3
     here, so the two differ by a sign. *)
  let a = add (pow x 3) (n 1) and b = add (pow x 3) x in
  check "resultant is antisymmetric in odd degrees"
    (equal (resultant "x" a b) (neg (resultant "x" b a)))

> resultant x^2 + a*x + b, x^2 + c*x + d, x
b^2 + d^2 - 2*b*d - a*b*c - a*c*d + b*c^2 + d*a^2
 
which is (b - d)^2 + (a - c)*(a*d - b*c), as it should be

(* Returns factors paired with their multiplicities, whose product is
   the primitive part of p, each factor square-free and the factors
   pairwise coprime. The trick is that p and p' share exactly the
   repeated factors, each with its multiplicity dropped by one, so one
   GCD peels off a whole layer at a time. Valid in characteristic
   zero only. *)
let square_free (p : t) (v : string) : (t * int) list =
  if is_zero p || degree_in p v <= 0 then [ (primitive p, 1) ]
  else begin
    let a = primitive p in
    let b = diff a v in
    let d = gcd a b in
    let w = ref (divide_exn a d) in
    let y = ref (divide_exn b d) in
    let z = ref (sub !y (diff !w v)) in
    let out = ref [] and i = ref 1 in
    while degree_in !w v > 0 do
      let g = gcd !w !z in
      if degree_in g v > 0 then out := (g, !i) :: !out;
      w := divide_exn !w g;
      y := divide_exn !z g;
      z := sub !y (diff !w v);
      incr i
    done;
    List.rev !out
  end

> sqfree (x - 1)^2*(x + 2)^3, x
(-1 + x)^2 * (2 + x)^3
> sqfree (x^2 - 1)^3*(x + 5)^2, x
(5 + x)^2 * (-1 + x^2)^3

exception Not_polynomial of string
 
type ctx = { mutable kernels : (string * Expr.t) list }
 
let new_ctx () : ctx = { kernels = [] }
 
let kernel (c : ctx) (e : Expr.t) : string =
  let name = Expr.to_string e in
  if not (List.mem_assoc name c.kernels) then c.kernels <- (name, e) :: c.kernels;
  name
 
(* ------------------------------------------------------------------ *)
(* Expr -> Poly                                                        *)
(* ------------------------------------------------------------------ *)
 
let rec to_poly (c : ctx) (e : Expr.t) : Poly.t =
  match e with
  | Expr.Num r -> Poly.of_rational r
  | Expr.Sym s ->
    ignore (kernel c (Expr.Sym s));
    Poly.var s
  | Expr.Add xs -> List.fold_left (fun acc x -> Poly.add acc (to_poly c x)) Poly.zero xs
  | Expr.Mul xs -> List.fold_left (fun acc x -> Poly.mul acc (to_poly c x)) Poly.one xs
  | Expr.Pow (b, Expr.Num n) when Rational.is_integer n && Rational.sign n >= 0 -> (
    match Rational.to_int_opt n with
    | Some k -> Poly.pow (to_poly c b) k
    | None -> raise (Not_polynomial "exponent too large to expand"))
  | Expr.Pow _ | Expr.Fun _ ->
    (* A negative or fractional power, or an unknown function: opaque. *)
    Poly.var (kernel c e)

let of_poly (c : ctx) (p : Poly.t) : Expr.t =
  let vars = Array.of_list (Poly.vars_of p) in
  let var_expr i =
    let name = vars.(i) in
    match List.assoc_opt name c.kernels with Some e -> e | None -> Expr.Sym name
  in
  let terms =
    List.map
      (fun (m, coeff) ->
        let factors = ref [] in
        Array.iteri
          (fun i k ->
            if k = 1 then factors := var_expr i :: !factors
            else if k > 1 then factors := Expr.Pow (var_expr i, Expr.int k) :: !factors)
          m;
        let body = match !factors with [] -> Expr.num_one | [ f ] -> f | fs -> Expr.Mul fs in
        if Rational.is_one coeff then body else Expr.Mul [ Expr.Num coeff; body ])
      (Poly.terms_of p)
  in
  match terms with [] -> Expr.num_zero | [ t ] -> Expr.simplify t | ts -> Expr.simplify (Expr.Add ts)
 
(* Round-trip an expression through the polynomial representation.
   This is `expand`: the polynomial form has no nesting left to
   distribute, so converting back is what performs the distribution. *)
let expand (e : Expr.t) : Expr.t =
  let c = new_ctx () in
  of_poly c (to_poly c (Expr.simplify e))

(* Round-trip an expression through the polynomial representation.
   This is `expand`: the polynomial form has no nesting left to
   distribute, so converting back is what performs the distribution. *)
let expand (e : Expr.t) : Expr.t =
  let c = new_ctx () in
  of_poly c (to_poly c (Expr.simplify e))

> expand (x + y)^3
x^3 + y^3 + 3*x*y^2 + 3*y*x^2
> expand (x + 1)*(x - 1)*(x^2 + 1)
-1 + x^4
> expand (x + sin(y))^2
x^2 + sin(y)^2 + 2*x*sin(y)
> expand (x^(1/2) + 1)^2
1 + x + 2*x^(1/2)
> expand (1/x + 1)^2
1 + x^(-2) + 2*x^(-1)

let degree (e : Expr.t) (v : string) : int =
  let c = new_ctx () in
  Poly.degree_in (to_poly c (Expr.simplify e)) v
 
let total_degree (e : Expr.t) : int =
  let c = new_ctx () in
  Poly.total_degree (to_poly c (Expr.simplify e))
 
(* The coefficient of v^k, as an expression. *)
let coeff (e : Expr.t) (v : string) (k : int) : Expr.t =
  let c = new_ctx () in
  of_poly c (Poly.coeff_in (to_poly c (Expr.simplify e)) v k)
 
(* Rewrite as a sum of coeff * v^k with the coefficients grouped. This
   is what `expand` is not: it keeps the structure the user asked for
   rather than flattening everything. *)
let collect (e : Expr.t) (v : string) : Expr.t =
  let c = new_ctx () in
  let p = to_poly c (Expr.simplify e) in
  let cs = Poly.coeffs_in p v in
  let x = Expr.Sym v in
  let terms =
    List.mapi (fun k ck -> (k, ck)) cs
    |> List.filter (fun (_, ck) -> not (Poly.is_zero ck))
    |> List.map (fun (k, ck) ->
         let ce = of_poly c ck in
         if k = 0 then ce
         else begin
           let xp = if k = 1 then x else Expr.Pow (x, Expr.int k) in
           if Expr.is_num_one ce then xp else Expr.Mul [ ce; xp ]
         end)
  in
  match terms with [] -> Expr.num_zero | [ t ] -> t | ts -> Expr.Add (List.rev ts)

> collect x^2*y + x^2 + 3*x*y + 7, x
(1 + y)*x^2 + 3*y*x + 7
> coeff (1 + x)^20, x, 10
184756
> degree (x*y + z)^7, z
7
> degree (x + y)^5 * (x - 1), x
6

exception Not_differentiable of string
 
(* Sums, products, and powers with a constant exponent - which is
   everything the polynomial layer can produce, plus the negative and
   fractional powers it treats as kernels. A power with a symbolic
   exponent needs a logarithm to differentiate, and a Fun node needs a
   derivative table; both wait for the calculus part. *)
let rec diff (e : Expr.t) (v : string) : Expr.t =
  Expr.simplify (diff_raw e v)
 
and diff_raw (e : Expr.t) (v : string) : Expr.t =
  match e with
  | Expr.Num _ -> Expr.num_zero
  | Expr.Sym s -> if s = v then Expr.num_one else Expr.num_zero
  | Expr.Add xs -> Expr.Add (List.map (fun x -> diff_raw x v) xs)
  | Expr.Mul xs ->
    (* Product rule, n-ary: one summand per factor differentiated. *)
    Expr.Add
      (List.mapi
         (fun i _ ->
           Expr.Mul (List.mapi (fun j x -> if i = j then diff_raw x v else x) xs))
         xs)
  | Expr.Pow (b, Expr.Num n) ->
    (* n * b^(n-1) * b' *)
    let n1 = Expr.Num (Rational.sub n Rational.one) in
    Expr.Mul [ Expr.Num n; Expr.Pow (b, n1); diff_raw b v ]
  | Expr.Pow _ -> raise (Not_differentiable "power with a symbolic exponent")
  | Expr.Fun (f, _) -> raise (Not_differentiable ("no derivative rule for " ^ f))

> diff (x^2 + 3*x + 1)^3, x
3*(1 + x^2 + 3*x)^2*(3 + 2*x)
> diff (x^3 + 1)^4, x
12*x^2*(1 + x^3)^3
> diff 1/x, x
-x^(-2)
> diff (x^2 + 1)^(1/2), x
x*(1 + x^2)^(-1/2)
> diff cos(x), x
cannot differentiate: no derivative rule for cos
> diff x^y, x
cannot differentiate: power with a symbolic exponent

let is_digit c = c >= '0' && c <= '9'
let is_alpha c = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c = '_'
let is_alnum c = is_alpha c || is_digit c
 
let tokenize (s : string) : token list =
  let n = String.length s in
  let out = ref [] in
  let i = ref 0 in
  while !i < n do
    let c = s.[!i] in
    if c = ' ' || c = '\t' || c = '\n' || c = '\r' then incr i
    else if is_digit c then begin
      let j = ref !i in
      while !j < n && is_digit s.[!j] do
        incr j
      done;
      out := TNum (Rational.of_string (String.sub s !i (!j - !i))) :: !out;
      i := !j
    end
    else if is_alpha c then begin
      let j = ref !i in
      while !j < n && is_alnum s.[!j] do
        incr j
      done;
      out := TIdent (String.sub s !i (!j - !i)) :: !out;
      i := !j
    end
    else begin
      let t =
        match c with
        | '+' -> TPlus
        | '-' -> TMinus
        | '*' -> TStar
        | '/' -> TSlash
        | '^' -> TCaret
        | '(' -> TLParen
        | ')' -> TRParen
        | ',' -> TComma
        | _ -> raise (Parse_error (Printf.sprintf "unexpected character %C" c))
      in
      out := t :: !out;
      incr i
    end
  done;
  List.rev (TEnd :: !out)

(* Left binding power. Higher binds tighter. A caret is given a lower
   right binding power than its left one when it recurses, which is
   what makes it right associative. *)
let lbp (t : token) : int =
  match t with
  | TPlus | TMinus -> 10
  | TStar | TSlash -> 20
  | TCaret -> 30
  | _ -> 0
 
let rec parse_expr (st : state) (bp : int) : Expr.t =
  let left = ref (parse_prefix st) in
  let continue_ = ref true in
  while !continue_ do
    let t = peek st in
    if lbp t <= bp then continue_ := false
    else begin
      ignore (advance st);
      left := parse_infix st !left t
    end
  done;
  !left
 
and parse_prefix (st : state) : Expr.t =
  match advance st with
  | TNum r -> Expr.Num r
  | TIdent name ->
    if peek st = TLParen then begin
      ignore (advance st);
      let args = parse_args st in
      Expr.Fun (name, args)
    end
    else Expr.Sym name
  (* Unary minus binds tighter than + and * but looser than ^, so that
     -x^2 parses as -(x^2) and -x*y as (-x)*y. *)
  | TMinus -> Expr.Mul [ Expr.Num Rational.minus_one; parse_expr st 25 ]
  | TPlus -> parse_expr st 25
  | TLParen ->
    let e = parse_expr st 0 in
    expect st TRParen ")";
    e
  | TEnd -> raise (Parse_error "unexpected end of input")
  | _ -> raise (Parse_error "expected an expression")

and parse_args (st : state) : Expr.t list =
  if peek st = TRParen then begin
    ignore (advance st);
    []
  end
  else begin
    let args = ref [ parse_expr st 0 ] in
    while peek st = TComma do
      ignore (advance st);
      args := parse_expr st 0 :: !args
    done;
    expect st TRParen ")";
    List.rev !args
  end
 
and parse_infix (st : state) (left : Expr.t) (t : token) : Expr.t =
  match t with
  | TPlus -> Expr.Add [ left; parse_expr st 10 ]
  | TMinus -> Expr.Add [ left; Expr.Mul [ Expr.Num Rational.minus_one; parse_expr st 10 ] ]
  | TStar -> Expr.Mul [ left; parse_expr st 20 ]
  | TSlash -> Expr.Mul [ left; Expr.Pow (parse_expr st 20, Expr.Num Rational.minus_one) ]
  | TCaret -> Expr.Pow (left, parse_expr st 29) (* 29 < 30: right associative *)
  | _ -> raise (Parse_error "not an infix operator")

let test_precedence () =
  check_tree "product binds tighter than sum" "(+ 2 (* 3 x))" "2 + 3*x";
  check_tree "power binds tighter than product" "(* 3 (^ x 2))" "3*x^2";
  check_tree "parentheses override" "(^ (+ x 1) 2)" "(x + 1)^2";
  check_tree "left associative subtraction" "(+ (+ a (* -1 b)) (* -1 c))" "a - b - c";
  check_tree "right associative power" "(^ 2 (^ 3 2))" "2^3^2";
  check_parse "and that means 512" "512" "2^3^2";
  check_tree "unary minus is looser than power" "(* -1 (^ x 2))" "-x^2";
  check_parse "so -x^2 at x=2 is -4" "-4" "-(2^2)";
  check_tree "unary minus is tighter than product" "(* (* -1 x) y)" "-x*y";
  check_tree "division is a negative power" "(* a (^ b -1))" "a/b";
  check_tree "chained division is left associative" "(* (* a (^ b -1)) (^ c -1))" "a/b/c"

(* Printing a simplified expression and parsing it back must land on
   the same expression. That is the check that the printer and the
   parser agree about precedence. *)
let test_round_trip () =
  let inputs =
    [
      "(x + y)^2";
      "1/2*x + 2/3*y";
      "x*y*z - 3";
      "x^2/(y + 1)";
      "-x - y";
      "2^x*3^y";
      "(x + 1)^(1/2)";
      "f(x, y)*g(z)";
      "x - y - z";
      "1/(x*y)";
    ]
  in
  List.iter
    (fun input ->
      let e = Parser.parse_simplified input in
      let printed = Expr.to_string e in
      incr checks;
      match Parser.parse_simplified printed with
      | exception Parser.Parse_error m ->
        incr failures;
        Printf.printf "FAIL: reparsing %s failed: %s\n" printed m
      | e2 ->
        if Expr.compare_expr e e2 <> 0 then begin
          incr failures;
          Printf.printf "FAIL: round trip changed %s\n  printed %s\n  reparsed %s\n" input printed
            (Expr.to_string e2)
        end)
    inputs;
  check "parse is total on its own output" true

let eval_command (line : string) : string =
  let line = String.trim line in
  let cmd, rest =
    match String.index_opt line ' ' with
    | None -> (line, "")
    | Some i -> (String.sub line 0 i, String.sub line (i + 1) (String.length line - i - 1))
  in
  let args = split_args rest in
  let arg n = List.nth args n in
  match (cmd, List.length args) with
  | "expand", 1 -> show (Algebra.expand (parse (arg 0)))
  | "simplify", 1 -> show (parse (arg 0))
  | "collect", 2 -> show (Algebra.collect (parse (arg 0)) (sym_name (arg 1)))
  | "degree", 2 -> string_of_int (Algebra.degree (parse (arg 0)) (sym_name (arg 1)))
  | "coeff", 3 ->
    show (Algebra.coeff (parse (arg 0)) (sym_name (arg 1)) (int_of_string (String.trim (arg 2))))
  | "diff", 2 -> show (Algebra.diff (parse (arg 0)) (sym_name (arg 1)))
  | "gcd", 2 -> show (Algebra.gcd (parse (arg 0)) (parse (arg 1)))
  | "lcm", 2 -> show (Algebra.lcm (parse (arg 0)) (parse (arg 1)))
  | "resultant", 3 ->
    show (Algebra.resultant (parse (arg 0)) (parse (arg 1)) (sym_name (arg 2)))
  | "sqfree", 2 ->
    Algebra.square_free (parse (arg 0)) (sym_name (arg 1))
    |> List.map (fun (f, m) -> Printf.sprintf "(%s)^%d" (show f) m)
    |> String.concat " * "
  | "tree", 1 -> Expr.to_sexp (parse (arg 0))
  | "help", _ ->
    "expand E | simplify E | collect E, v | degree E, v | coeff E, v, k\n\
     diff E, v | gcd E, F | lcm E, F | resultant E, F, v | sqfree E, v | tree E"
  (* Anything that is not a command word is treated as an expression to
     simplify, so the REPL can be used as a calculator. *)
  | _ -> show (parse line)

(* Split on commas that are not inside parentheses, so that
   "coeff (x+1)^3, x, 2" separates into three arguments. *)
let split_args (s : string) : string list =
  let out = ref [] and buf = Buffer.create 32 and depth = ref 0 in
  String.iter
    (fun c ->
      if c = '(' then begin
        incr depth;
        Buffer.add_char buf c
      end
      else if c = ')' then begin
        decr depth;
        Buffer.add_char buf c
      end
      else if c = ',' && !depth = 0 then begin
        out := Buffer.contents buf :: !out;
        Buffer.clear buf
      end
      else Buffer.add_char buf c)
    s;
  out := Buffer.contents buf :: !out;
  List.rev_map String.trim !out

$ ./_out/main repl
cas part 3 - type `help`, or an expression to simplify
> 2^64 + 1
18446744073709551617
> 1/3 + 1/6
1/2
> expand (x + y)^3
x^3 + y^3 + 3*x*y^2 + 3*y*x^2
> gcd x^6 - 1, x^4 - 1
-1 + x^2
> resultant x^2 + a*x + b, x^2 + c*x + d, x
b^2 + d^2 - 2*b*d - a*b*c - a*c*d + b*c^2 + d*a^2
> sqfree (x^2 - 1)^3*(x + 5)^2, x
(5 + x)^2 * (-1 + x^2)^3
> coeff (1 + x)^20, x, 10
184756
> foo x
parse error: trailing input
> quit

(* The property that matters: whatever the two algorithms return, it
   must divide both inputs, and the quotients must be coprime. That is
   the definition of a GCD, and it is checkable without knowing the
   answer in advance. *)
let random_poly deg =
  Upoly.of_list (List.init (deg + 1) (fun i -> if i = deg then 1 + Random.int 9 else Random.int 21 - 10))
 
let test_gcd_property () =
  Random.init 20260420;
  for _ = 1 to 200 do
    let g = random_poly (1 + Random.int 4) in
    let a = Upoly.mul g (random_poly (1 + Random.int 4)) in
    let b = Upoly.mul g (random_poly (1 + Random.int 4)) in
    let d1 = Upoly.gcd_prs a b and d2 = Upoly.gcd_modular a b in
    check "the two algorithms agree" (Upoly.equal d1 d2);
    check "the gcd divides a" (Upoly.divides d1 a <> None);
    check "the gcd divides b" (Upoly.divides d1 b <> None);
    check "the gcd is a multiple of the planted factor"
      (Upoly.divides (Upoly.primitive_part g) d1 <> None);
    match (Upoly.divides d1 a, Upoly.divides d1 b) with
    | Some qa, Some qb ->
      check "the cofactors are coprime" (Upoly.degree (Upoly.gcd_prs qa qb) = 0)
    | _ -> ()
  done

let test_ring_laws () =
  Random.init 424242;
  for _ = 1 to 300 do
    let a = random_poly 2 and b = random_poly 2 and c = random_poly 2 in
    check "addition commutes" (equal (add a b) (add b a));
    check "multiplication commutes" (equal (mul a b) (mul b a));
    check "multiplication associates" (equal (mul (mul a b) c) (mul a (mul b c)));
    check "multiplication distributes" (equal (mul a (add b c)) (add (mul a b) (mul a c)));
    check "subtraction inverts addition" (equal (sub (add a b) b) a);
    (* A product always divides back exactly. *)
    if not (is_zero b) then
      check "exact division undoes multiplication" (equal (divide_exn (mul a b) b) a);
    (* Degrees add under multiplication. *)
    if (not (is_zero a)) && not (is_zero b) then
      check "degrees add" (total_degree (mul a b) = total_degree a + total_degree b)
  done

(* The property that catches an expansion bug immediately: expanding
   must not change what an expression evaluates to. *)
let test_expand_preserves_value () =
  Random.init 31415;
  let rec random_expr d =
    if d <= 0 then
      match Random.int 4 with
      | 0 -> Expr.int (Random.int 5 - 2)
      | 1 -> Expr.rat (Random.int 5 - 2) (1 + Random.int 3)
      | 2 -> Expr.Sym "x"
      | _ -> Expr.Sym "y"
    else
      match Random.int 4 with
      | 0 -> Expr.Add [ random_expr (d - 1); random_expr (d - 1) ]
      | 1 -> Expr.Mul [ random_expr (d - 1); random_expr (d - 1) ]
      | 2 -> Expr.Add [ random_expr (d - 1); random_expr (d - 1); random_expr (d - 1) ]
      | _ -> Expr.Pow (random_expr (d - 1), Expr.int (Random.int 4))
  in
  for _ = 1 to 400 do
    let e = random_expr 3 in
    let xv = Expr.int (1 + Random.int 4) and yv = Expr.int (1 + Random.int 4) in
    let value g = Expr.to_rational (Expr.substitute "y" yv (Expr.substitute "x" xv g)) in
    match (value e, value (Algebra.expand e)) with
    | exception Expr.Undefined _ -> ()
    | exception Division_by_zero -> ()
    | Some a, Some b ->
      check "expansion preserves the value" (Rational.equal a b)
    | _ -> ()
  done

no factorization        expand ((x+1)*(x+2)) gives x^2 + 3x + 2, and nothing
                        turns it back. square_free splits multiplicities,
                        not factors.
 
no rational functions   1/(x^2 - 1) + 1/(x + 1) stays as written. Adding
                        them needs a common denominator and cancelling
                        the result needs the GCD this part just built.
 
Part 4: factorization over F_p, Hensel lifting to Z, multivariate
        factorization, and the rational function field on top of it.