Building a CAS in OCaml Part 4

f over Z
  |  square-free split            Yun, from Part 3
  |  make monic                   change of variable
  v
f mod p                           pick p keeping degree and squarefreeness
  |  square-free mod p            Musser, because x^p has zero derivative
  |  distinct-degree              gcd with x^(p^d) - x
  |  equal-degree                 Cantor-Zassenhaus, randomized
  v
u_1 * ... * u_r mod p
  |  Hensel lift                  p -> p^k, k from a coefficient bound
  v
u_1 * ... * u_r mod p^k
  |  recombine                    which subsets divide f over Z
  v
irreducible factors over Z

module Zp = struct
  let norm p a =
    let r = a mod p in
    if r < 0 then r + p else r
 
  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
 
  let pow p a n =
    let rec go acc b n =
      if n = 0 then acc else if n land 1 = 1 then go (mul p acc b) (mul p b b) (n lsr 1)
      else go acc (mul p b b) (n lsr 1)
    in
    go 1 (norm p a) n
 
  (* Extended Euclid, iterative, tracking only the coefficient needed. *)
  let inv p a =
    let a = norm p a in
    if a = 0 then invalid_arg "Fp.Zp.inv: zero has no inverse";
    let r0 = ref p and r1 = ref a 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;
    norm p !t0
end

(* Division with remainder. F_p is a field, so unlike Z there is no
   pseudo-division here: the leading coefficient is always invertible. *)
let divmod p (a : t) (b : t) : t * t =
  if is_zero b then raise Division_by_zero;
  let db = degree b and ib = Zp.inv p (lc b) in
  if degree a < db then (zero, a)
  else begin
    let q = Array.make (degree a - db + 1) 0 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
      q.(d) <- c;
      r := sub p !r (shift (scale p b c) d)
    done;
    (normalize q, !r)
  end

(* Extended Euclid: returns (g, s, t) with s*a + t*b = g, g monic.
   Hensel lifting needs the cofactors, not just the GCD. *)
let ext_gcd p (a : t) (b : t) : t * t * t =
  let r0 = ref a and r1 = ref b in
  let s0 = ref one and s1 = ref zero in
  let t0 = ref zero and t1 = ref one in
  while not (is_zero !r1) do
    let q, r = divmod p !r0 !r1 in
    r0 := !r1;
    r1 := r;
    let s = sub p !s0 (mul p q !s1) in
    s0 := !s1;
    s1 := s;
    let t = sub p !t0 (mul p q !t1) in
    t0 := !t1;
    t1 := t
  done;
  let c = Zp.inv p (lc !r0) in
  (scale p !r0 c, scale p !s0 c, scale p !t0 c)

(* a^n mod m, by repeated squaring in the quotient ring F_p[x]/(m). *)
let pow_mod p (a : t) (n : int) (m : t) : t =
  let rec go acc b n =
    if n = 0 then acc
    else if n land 1 = 1 then go (rem p (mul p acc b) m) (rem p (mul p b b) m) (n lsr 1)
    else go acc (rem p (mul p b b) m) (n lsr 1)
  in
  go one (rem p a m) n
 
(* The same with an exponent too large for an int: equal-degree
   factorization needs a^((p^d - 1)/2), which overflows immediately. *)
let pow_mod_big p (a : t) (n : Bigint.t) (m : t) : t =
  let two = Bigint.of_int 2 in
  let acc = ref one and b = ref (rem p a m) and e = ref n in
  while not (Bigint.is_zero !e) do
    let q, r = Bigint.divmod !e two in
    if not (Bigint.is_zero r) then acc := rem p (mul p !acc !b) m;
    b := rem p (mul p !b !b) m;
    e := q
  done;
  !acc

let pth_root p (a : t) : t =
  let n = degree a in
  normalize (Array.init ((n / p) + 1) (fun i -> a.(i * p)))
 
let rec square_free p (f : t) : (t * int) list =
  if degree f <= 0 then []
  else begin
    let f = monic p f in
    let fd = diff p f in
    if is_zero fd then
      (* f is a p-th power: f(x) = g(x)^p, so every multiplicity in g
         is multiplied by p. *)
      List.map (fun (g, m) -> (g, m * p)) (square_free p (pth_root p f))
    else begin
      let c = ref (gcd p f fd) in
      let w = ref (div p f !c) in
      let out = ref [] and i = ref 1 in
      while degree !w > 0 do
        let y = gcd p !w !c in
        let z = div p !w y in
        if degree z > 0 then out := (z, !i) :: !out;
        w := y;
        c := div p !c y;
        incr i
      done;
      (* Whatever is left in c is a p-th power. *)
      let tail =
        if degree !c > 0 then
          List.map (fun (g, m) -> (g, m * p)) (square_free p (pth_root p !c))
        else []
      in
      List.rev_append !out tail
    end
  end

  (* The characteristic-p case that Yun cannot see: (x+1)^13 has a zero
     derivative, so the algorithm has to take a p-th root instead. *)
  let h = Fp.of_list p (List.init 14 (fun i -> if i = 0 || i = 13 then 1 else 0)) in
  check_str "(x+1)^13 is x^13 + 1 mod 13" "x^13 + 1" (s h);
  check "derivative vanishes" (Fp.is_zero (Fp.diff p h));
  match Fp.square_free p h with
  | [ (g, 13) ] -> check "p-th power decomposes correctly" (Fp.equal g (f [ 1; 1 ]))
  | l ->
    incr checks;
    incr failures;
    Printf.printf "FAIL: p-th power gave %d pieces\n" (List.length l)

let distinct_degree p (f : t) : (t * int) list =
  let out = ref [] in
  let fstar = ref f in
  let xp = ref x in
  let d = ref 1 in
  while degree !fstar >= 2 * !d do
    xp := pow_mod p !xp p !fstar;
    let g = gcd p !fstar (sub p !xp x) in
    if degree g > 0 then begin
      out := (g, !d) :: !out;
      fstar := div p !fstar g;
      xp := rem p !xp !fstar
    end;
    incr d
  done;
  (* Anything left is irreducible: its degree exceeds half the degree
     of what remained, so it cannot be a product of two pieces. *)
  if degree !fstar > 0 then out := (!fstar, degree !fstar) :: !out;
  List.rev !out

let rec equal_degree p (f : t) (d : int) : t list =
  if degree f = 0 then []
  else if degree f = d then [ monic p f ]
  else begin
    let n = degree f in
    let e =
      (* (p^d - 1) / 2, which needs a bignum for any interesting d. *)
      Bigint.div (Bigint.sub (Bigint.pow (Bigint.of_int p) d) Bigint.one) (Bigint.of_int 2)
    in
    let result = ref None in
    while !result = None do
      let a = random p (n - 1) in
      if degree a > 0 then begin
        let g = gcd p a f in
        let split =
          if degree g > 0 && degree g < n then Some g
          else begin
            let b = pow_mod_big p a e f in
            let g = gcd p (sub p b one) f in
            if degree g > 0 && degree g < n then Some g else None
          end
        in
        match split with
        | Some g -> result := Some (equal_degree p g d @ equal_degree p (div p f g) d)
        | None -> ()
      end
    done;
    match !result with Some l -> l | None -> assert false
  end

let factor p (f : t) : int * (t * int) list =
  if is_zero f then invalid_arg "Fp.factor: zero";
  let c = lc f in
  let f = monic p f in
  let out = ref [] in
  List.iter
    (fun (g, m) ->
      List.iter
        (fun (h, d) -> List.iter (fun q -> out := (q, m) :: !out) (equal_degree p h d))
        (distinct_degree p g))
    (square_free p f);
  (c, List.sort (fun (a, _) (b, _) -> compare a b) !out)
 
let is_irreducible p (f : t) : bool =
  degree f > 0 && (match factor p f with _, [ (_, 1) ] -> true | _ -> false)

let berlekamp_matrix p (f : t) : int array array =
  let n = degree f in
  let q = Array.make_matrix n n 0 in
  for i = 0 to n - 1 do
    let row = pow_mod p (shift one i) p f in
    for j = 0 to n - 1 do
      q.(j).(i) <- coeff row j
    done;
    (* Subtract the identity: the matrix of v -> v^p - v. *)
    q.(i).(i) <- Zp.sub p q.(i).(i) 1
  done;
  q

(* Gaussian elimination over F_p, returning a basis of the nullspace. *)
let nullspace p (m : int array array) : int array list =
  let rows = Array.length m in
  if rows = 0 then []
  else begin
    let cols = Array.length m.(0) in
    let a = Array.map Array.copy m in
    let pivot_of_col = Array.make cols (-1) in
    let r = ref 0 in
    for c = 0 to cols - 1 do
      if !r < rows then begin
        (* Find a row with a nonzero entry in this column. *)
        let piv = ref (-1) in
        for i = rows - 1 downto !r do
          if a.(i).(c) <> 0 then piv := i
        done;
        if !piv >= 0 then begin
          let tmp = a.(!r) in
          a.(!r) <- a.(!piv);
          a.(!piv) <- tmp;
          let inv = Zp.inv p a.(!r).(c) in
          for j = 0 to cols - 1 do
            a.(!r).(j) <- Zp.mul p a.(!r).(j) inv
          done;
          for i = 0 to rows - 1 do
            if i <> !r && a.(i).(c) <> 0 then begin
              let factor = a.(i).(c) in
              for j = 0 to cols - 1 do
                a.(i).(j) <- Zp.sub p a.(i).(j) (Zp.mul p factor a.(!r).(j))
              done
            end
          done;
          pivot_of_col.(c) <- !r;
          incr r
        end
      end
    done;
    (* One basis vector per free column. *)
    let basis = ref [] in
    for c = cols - 1 downto 0 do
      if pivot_of_col.(c) < 0 then begin
        let v = Array.make cols 0 in
        v.(c) <- 1;
        for c2 = 0 to cols - 1 do
          if pivot_of_col.(c2) >= 0 then v.(c2) <- Zp.neg p a.(pivot_of_col.(c2)).(c)
        done;
        basis := v :: !basis
      end
    done;
    !basis
  end

(* Berlekamp shares none of the code of distinct-degree plus
   equal-degree, so its factor count is an independent check on both. *)
let test_berlekamp_agrees () =
  Random.init 20260427;
  let primes = [ 5; 7; 11; 13; 17 ] in
  List.iter
    (fun q ->
      for _ = 1 to 40 do
        let d = 2 + Random.int 5 in
        let g = Fp.monic q (Fp.random q d) in
        if Fp.degree g = d && Fp.degree (Fp.gcd q g (Fp.diff q g)) = 0 then begin
          let _, fs = Fp.factor q g in
          check "berlekamp nullity equals the factor count"
            (Fp.berlekamp_count q g = List.length fs);
          check "the factors multiply back"
            (Fp.equal (List.fold_left (fun acc (h, _) -> Fp.mul q acc h) Fp.one fs) g);
          check "every factor is irreducible"
            (List.for_all (fun (h, _) -> Fp.is_irreducible q h) fs)
        end
      done)
    primes

(* Every factor of f has coefficients bounded by 2^n times the
   2-norm of f, and the 2-norm is at most (n+1) times the max norm.
   Generous, and cheap to compute, and what matters is that it is an
   actual bound and not an estimate. *)
let coefficient_bound (f : Upoly.t) : Bigint.t =
  let n = Upoly.degree f in
  Bigint.mul (Bigint.pow (Bigint.of_int 2) n)
    (Bigint.mul (Bigint.of_int (n + 1)) (Upoly.max_norm f))
 
(* Smallest k with p^k > 2*bound, so that a symmetric representative
   modulo p^k is the integer itself rather than a congruent stand-in. *)
let lift_exponent (p : int) (bound : Bigint.t) : int =
  let target = Bigint.mul (Bigint.of_int 2) bound in
  let k = ref 1 and m = ref (Bigint.of_int p) in
  while Bigint.compare !m target <= 0 do
    m := Bigint.mul !m (Bigint.of_int p);
    incr k
  done;
  !k

(* Coefficients reduced into (-m/2, m/2]. Keeping the symmetric
   representative throughout means the lifted factors are already the
   integer factors once the modulus is large enough, with no final
   adjustment. *)
let center (a : Upoly.t) (m : Bigint.t) : Upoly.t =
  let half = Bigint.div m (Bigint.of_int 2) in
  Upoly.normalize
    (Array.map
       (fun c ->
         let r = Bigint.rem c m in
         let r = if Bigint.sign r < 0 then Bigint.add r m else r in
         if Bigint.compare r half > 0 then Bigint.sub r m else r)
       a)

let test_bound () =
  (* The bound has to be a bound, not an estimate: every coefficient of
     every factor must fit inside it. *)
  let f = p [ -6; 11; -6; 1 ] in
  let b = Factor.coefficient_bound f in
  let _, fs = Factor.factor f in
  check "every factor fits inside the bound"
    (List.for_all (fun (g, _) -> Bigint.compare (Upoly.max_norm g) b <= 0) fs);
  let k = Factor.lift_exponent 5 b in
  check "the lift exponent clears twice the bound"
    (Bigint.compare (Bigint.pow (Bigint.of_int 5) k) (Bigint.mul (Bigint.of_int 2) b) > 0);
  check "and is minimal"
    (Bigint.compare (Bigint.pow (Bigint.of_int 5) (k - 1)) (Bigint.mul (Bigint.of_int 2) b) <= 0)

let hensel_step (p : int) (m : Bigint.t) (f : Upoly.t) (g : Upoly.t) (h : Upoly.t)
    (g1 : Fp.t) (h1 : Fp.t) (s : Fp.t) (t : Fp.t) : Upoly.t * Upoly.t =
  let err = Upoly.sub f (Upoly.mul g h) in
  let c = Upoly.divexact_int err m in
  let cp = Upoly.to_fp c p in
  let u0 = Fp.rem p (Fp.mul p t cp) g1 in
  let q = Fp.div p (Fp.mul p t cp) g1 in
  let v0 = Fp.add p (Fp.mul p s cp) (Fp.mul p q h1) in
  let u = Upoly.of_fp u0 p and v = Upoly.of_fp v0 p in
  (Upoly.add g (Upoly.scale u m), Upoly.add h (Upoly.scale v m))

(* Lift a two-factor split all the way to p^k. *)
let hensel_lift2 (p : int) (k : int) (f : Upoly.t) (g1 : Fp.t) (h1 : Fp.t) : Upoly.t * Upoly.t =
  let d, s, t = Fp.ext_gcd p g1 h1 in
  if not (Fp.is_one d) then invalid_arg "hensel_lift2: factors are not coprime mod p";
  let g = ref (Upoly.of_fp g1 p) and h = ref (Upoly.of_fp h1 p) in
  let m = ref (Bigint.of_int p) in
  for _ = 1 to k - 1 do
    let g', h' = hensel_step p !m f !g !h g1 h1 s t in
    m := Bigint.mul !m (Bigint.of_int p);
    g := center g' !m;
    h := center h' !m
  done;
  (!g, !h)
 
(* Lift a factorization into any number of pieces, by splitting off one
   at a time. f must be monic, and every factor monic mod p. *)
let rec hensel_lift_many (p : int) (k : int) (f : Upoly.t) (factors : Fp.t list) : Upoly.t list =
  match factors with
  | [] -> []
  | [ _ ] -> [ f ]
  | g1 :: rest ->
    let h1 = List.fold_left (fun acc q -> Fp.mul p acc q) Fp.one rest in
    let g, h = hensel_lift2 p k f g1 h1 in
    g :: hensel_lift_many p k h rest

(* A usable prime has to leave the degree alone and leave the input
   square-free. The second condition is the one that matters: a prime
   dividing the discriminant merges two distinct factors into one and
   the lift would then be lifting the wrong factorization. *)
let choose_prime (f : Upoly.t) : int =
  (* No prime can work if f is not square-free over Z to begin with:
     the gcd condition below would never hold, and the search would run
     forever. Refuse rather than spin. *)
  if Upoly.degree (Upoly.gcd_modular f (Upoly.diff f)) > 0 then raise Not_square_free;
  let lead = Upoly.lc f in
  let p = ref 3 and found = ref 0 in
  while !found = 0 do
    p := Fp.next_prime (!p + 1);
    let pi = !p in
    if Upoly.bigint_mod_int lead pi <> 0 then begin
      let fp = Upoly.to_fp f pi in
      if Fp.degree fp = Upoly.degree f && Fp.degree (Fp.gcd pi fp (Fp.diff pi fp)) = 0 then
        found := pi
    end
  done;
  !found

(* The Hensel step is an identity, not a heuristic: after lifting to
   p^k the product must be congruent to f modulo p^k exactly. *)
let test_hensel () =
  let f = p [ 2; 3; 1 ] in
  (* x^2 + 3x + 2 = (x+1)(x+2), and mod 5 the factors are coprime *)
  let q = 5 in
  let g1 = Fp.of_list q [ 1; 1 ] and h1 = Fp.of_list q [ 2; 1 ] in
  List.iter
    (fun k ->
      let g, h = Factor.hensel_lift2 q k f g1 h1 in
      let m = Bigint.pow (Bigint.of_int q) k in
      let err = Upoly.sub f (Upoly.mul g h) in
      check (Printf.sprintf "lift to 5^%d is congruent" k)
        (Array.for_all (fun c -> Bigint.is_zero (Bigint.rem c m)) err);
      check (Printf.sprintf "lift to 5^%d keeps the degrees" k)
        (Upoly.degree g = 1 && Upoly.degree h = 1))
    [ 1; 2; 3; 5; 8 ];
  (* Lifted far enough, the factors are the integer ones. *)
  let k = Factor.lift_exponent q (Factor.coefficient_bound f) in
  let g, h = Factor.hensel_lift2 q k f g1 h1 in
  check "far enough is exact over Z" (Upoly.equal (Upoly.mul g h) f);
  (* A non-trivial case, where the true factors have coefficients
     larger than the prime. *)
  let f2 = Upoly.mul (p [ 37; 1 ]) (p [ -41; 1 ]) in
  let q2 = 7 in
  let a1 = Upoly.to_fp (p [ 37; 1 ]) q2 and b1 = Upoly.to_fp (p [ -41; 1 ]) q2 in
  let k2 = Factor.lift_exponent q2 (Factor.coefficient_bound f2) in
  let g2, h2 = Factor.hensel_lift2 q2 k2 f2 a1 b1 in
  check "coefficients larger than the prime are recovered"
    (Upoly.equal (Upoly.mul g2 h2) f2)

(* Enumerate the subsets of a list by size, smallest first, because the
   true factors are usually products of few lifted pieces and stopping
   early is the whole game. *)
let subsets_of_size (k : int) (l : 'a list) : 'a list list =
  let rec go k l =
    if k = 0 then [ [] ]
    else
      match l with
      | [] -> []
      | x :: rest -> List.map (fun s -> x :: s) (go (k - 1) rest) @ go k rest
  in
  go k l
 
(* Try products of the lifted factors against f. A subset whose product
   divides f exactly over Z is a true factor, and the pieces it used
   are removed before the search continues. *)
let recombine (f : Upoly.t) (lifted : Upoly.t list) (m : Bigint.t) : Upoly.t list =
  let remaining = ref lifted in
  let current = ref f in
  let out = ref [] in
  let size = ref 1 in
  while List.length !remaining > 0 && !size <= List.length !remaining / 2 do
    let found = ref false in
    let candidates = subsets_of_size !size !remaining in
    List.iter
      (fun subset ->
        if not !found then begin
          let prod = center (List.fold_left Upoly.mul Upoly.one subset) m in
          if Upoly.degree prod > 0 then
            match Upoly.divides (Upoly.primitive_part prod) !current with
            | Some q ->
              found := true;
              out := Upoly.primitive_part prod :: !out;
              current := q;
              remaining := List.filter (fun g -> not (List.memq g subset)) !remaining
            | None -> ()
        end)
      candidates;
    if not !found then incr size
  done;
  if Upoly.degree !current > 0 then out := !current :: !out;
  List.rev !out

(* Input must be monic, square-free, of positive degree. *)
let factor_monic_squarefree (f : Upoly.t) : Upoly.t list =
  if Upoly.degree f <= 1 then [ f ]
  else begin
    let p = choose_prime f in
    let fp = Fp.monic p (Upoly.to_fp f p) in
    let _, mods = Fp.factor p fp in
    let mods = List.map fst mods in
    if List.length mods = 1 then [ f ]
    else begin
      let k = lift_exponent p (coefficient_bound f) in
      let m = Bigint.pow (Bigint.of_int p) k in
      let lifted = hensel_lift_many p k f mods in
      recombine f lifted m
    end
  end

(* A non-monic f is made monic by a change of variable rather than by
   dragging the leading coefficient through the lift: with b = lc(f) and
   n = deg(f), the polynomial b^(n-1) * f(x/b) is monic with integer
   coefficients, and a factor g(x) of it maps back to the primitive
   part of g(b*x). *)
let to_monic (f : Upoly.t) : Upoly.t =
  let n = Upoly.degree f in
  let b = Upoly.lc f in
  Upoly.normalize
    (Array.init (n + 1) (fun i ->
       (* coefficient of x^i is a_i * b^(n-1-i) *)
       if i = n then Bigint.one
       else Bigint.mul (Upoly.coeff f i) (Bigint.pow b (n - 1 - i))))
 
let from_monic (g : Upoly.t) (b : Bigint.t) : Upoly.t =
  (* g(b*x), then take the primitive part. *)
  Upoly.primitive_part
    (Upoly.normalize
       (Array.init (Upoly.degree g + 1) (fun i -> Bigint.mul (Upoly.coeff g i) (Bigint.pow b i))))
 
let factor_squarefree (f : Upoly.t) : Upoly.t list =
  if Upoly.degree f <= 0 then []
  else begin
    let f = Upoly.primitive_part f in
    if Upoly.degree f <= 1 then [ f ]
    else begin
      let b = Upoly.lc f in
      if Bigint.equal b Bigint.one then factor_monic_squarefree f
      else
        List.map (fun g -> from_monic g b) (factor_monic_squarefree (to_monic f))
    end
  end

(* The whole thing: content, then square-free split, then factor each
   square-free piece. Returns the integer content and the irreducible
   factors with their multiplicities. *)
let factor (f : Upoly.t) : Bigint.t * (Upoly.t * int) list =
  if Upoly.is_zero f then (Bigint.zero, [])
  else begin
    let c = Upoly.content f in
    let c = if Bigint.sign (Upoly.lc f) < 0 then Bigint.neg c else c in
    let prim = Upoly.divexact_int f c in
    if Upoly.degree prim = 0 then (c, [])
    else begin
      let out = ref [] in
      List.iter
        (fun (g, m) -> List.iter (fun q -> out := (q, m) :: !out) (factor_squarefree g))
        (square_free prim);
      (c, List.rev !out)
    end
  end

> factor x^4 - 1, x
(1 + x)*(-1 + x)*(1 + x^2)
> factor 6*x^6 - 6, x
6*(1 + x)*(-1 + x)*(1 + x + x^2)*(1 + x^2 - x)
> factor 9*x^2 + 12*x + 4, x
(2 + 3*x)^2
> factor x^16 - 1, x
(1 + x)*(-1 + x)*(1 + x^2)*(1 + x^4)*(1 + x^8)

> factor x^4 + 1, x
1 + x^4
> factormod x^4 + 1, 5
1 * (x^2 + 2)^1 * (x^2 + 3)^1
> factormod x^4 + 1, 13
1 * (x^2 + 5)^1 * (x^2 + 8)^1

let swinnerton_dyer (k : int) : Upoly.t =
  (* The product over every sign choice of x - (+-sqrt p1 +- sqrt p2 ...).
     Adjoining one square root at a time: if f has roots r_i, then
     f(x - sqrt q) * f(x + sqrt q) has roots r_i +- sqrt q. Writing
     f(x + sqrt q) = A(x) + sqrt q * B(x) makes f(x - sqrt q) its
     conjugate, so the product is the norm A^2 - q*B^2 and stays in
     Z[x]. A and B come out of a Horner pass in Z[sqrt q][x]. *)
  let adjoin (f : Upoly.t) (q : int) : Upoly.t =
    let qb = Bigint.of_int q in
    let x = Upoly.of_list [ 0; 1 ] in
    (* (a, b) stands for a + sqrt q * b; multiplying by (x + sqrt q)
       gives (a*x + q*b, a + b*x). *)
    let a = ref Upoly.zero and b = ref Upoly.zero in
    for i = Upoly.degree f downto 0 do
      let a' = Upoly.add (Upoly.mul !a x) (Upoly.scale !b qb) in
      let b' = Upoly.add !a (Upoly.mul !b x) in
      a := Upoly.add a' (Upoly.const (Upoly.coeff f i));
      b := b'
    done;
    Upoly.sub (Upoly.mul !a !a) (Upoly.scale (Upoly.mul !b !b) qb)
  in
  let primes = [ 2; 3; 5; 7 ] in
  List.fold_left adjoin (Upoly.of_list [ 0; 1 ]) (List.filteri (fun i _ -> i < k) primes)

./_out/main factorbench

polynomial                         degree    mod p   over Z   time (s)
x^16 - 1                               16        8        5     0.0010
x^32 - 1                               32       10        6     0.0029
x^64 - 1                               64       12        7     0.0169
Swinnerton-Dyer 2 (sqrt 2, 3)           4        2        1     0.0001
Swinnerton-Dyer 3 (+ sqrt 5)            8        4        1     0.0003
Swinnerton-Dyer 4 (+ sqrt 7)           16        8        1     0.0049
product of 5 random irreducibles       25        8        5     0.0049

type t = {
  num : Poly.t;
  den : Poly.t; (* never zero, coprime with num, positive lex-leading coefficient *)
}
 
(* Cancel the common factor and fix the sign. Dividing by the GCD is
   what makes this a normal form, and it is also the single most
   expensive thing a rational function arithmetic does, which is why
   Part 3 spent so long making the GCD fast. *)
let make (num : Poly.t) (den : Poly.t) : t =
  if Poly.is_zero den then raise Division_by_zero;
  if Poly.is_zero num then { num = Poly.zero; den = Poly.one }
  else begin
    let g = Poly.gcd num den in
    let n = Poly.divide_exn num g and d = Poly.divide_exn den g in
    (* Push the content and the sign into the numerator so that the
       denominator is primitive with a positive leading coefficient. *)
    let dp = Poly.primitive d in
    let scale =
      match (Poly.lead_term d, Poly.lead_term dp) with
      | Some (_, c1), Some (_, c2) -> Rational.div c1 c2
      | _ -> Rational.one
    in
    { num = Poly.scale n (Rational.inv scale); den = dp }
  end

let neg (r : t) : t = { r with num = Poly.neg r.num }
 
let add (a : t) (b : t) : t =
  make (Poly.add (Poly.mul a.num b.den) (Poly.mul b.num a.den)) (Poly.mul a.den b.den)
 
let sub (a : t) (b : t) : t = add a (neg b)
let mul (a : t) (b : t) : t = make (Poly.mul a.num b.num) (Poly.mul a.den b.den)
 
let inv (a : t) : t =
  if is_zero a then raise Division_by_zero else make a.den a.num
 
let div (a : t) (b : t) : t =
  if is_zero b then raise Division_by_zero else make (Poly.mul a.num b.den) (Poly.mul a.den b.num)
 
let pow (a : t) (n : int) : t =
  if n >= 0 then make (Poly.pow a.num n) (Poly.pow a.den n)
  else inv (make (Poly.pow a.num (-n)) (Poly.pow a.den (-n)))
 
let equal (a : t) (b : t) : bool = Poly.equal a.num b.num && Poly.equal a.den b.den
 
(* d/dv of n/d is (n'd - nd') / d^2, and the normalizing constructor
   cancels whatever that leaves behind. *)
let diff (a : t) (v : string) : t =
  make
    (Poly.sub (Poly.mul (Poly.diff a.num v) a.den) (Poly.mul a.num (Poly.diff a.den v)))
    (Poly.mul a.den a.den)

> cancel (x^2 - 1)/(x + 1)
-1 + x
> together 1/(x + 1) + 1/(x - 1)
2*x*(-1 + x^2)^(-1)
> together 1/(x - 1) - 1/x - 1/(x*(x - 1))
0

(* Ordinary division with remainder, which needs the leading
   coefficient to be invertible. Over Q that means a nonzero constant,
   so this is the univariate case and nothing else; everywhere else
   pseudo-division is the only option. *)
let quo_rem (v : string) (a : t) (b : t) : t * t =
  if is_zero b then raise Division_by_zero;
  let lcb = lc_in b v in
  let inv =
    match to_rational lcb with
    | Some r when not (Rational.is_zero r) -> Rational.inv r
    | _ -> raise (Not_univariate "quo_rem: leading coefficient is not a constant")
  in
  let db = degree_in b v in
  let x = var v 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 = scale (lc_in !r v) inv in
    let term = mul c (pow x d) in
    q := add !q term;
    r := sub !r (mul term b)
  done;
  (!q, !r)

(* Extended Euclid over Q[v]: returns (g, s, t) with s*a + t*b = g. The
   partial-fraction split is exactly this identity applied to two
   coprime denominators. *)
let ext_gcd (v : string) (a : t) (b : t) : t * t * t =
  let r0 = ref a and r1 = ref b in
  let s0 = ref one and s1 = ref zero in
  let t0 = ref zero and t1 = ref one in
  while not (is_zero !r1) do
    let q, r = quo_rem v !r0 !r1 in
    r0 := !r1;
    r1 := r;
    let s = sub !s0 (mul q !s1) in
    s0 := !s1;
    s1 := s;
    let t = sub !t0 (mul q !t1) in
    t0 := !t1;
    t1 := t
  done;
  match to_rational (lc_in !r0 v) with
  | Some c when not (Rational.is_zero c) ->
    let i = Rational.inv c in
    (scale !r0 i, scale !s0 i, scale !t0 i)
  | _ -> (!r0, !s0, !t0)
 
(* ------------------------------------------------------------------ *)
(* Integer normalization                                               *)
(* ------------------------------------------------------------------ *)
 
(* 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

(* Splitting n/(d1*d2) with d1 and d2 coprime is the Bezout identity
   and nothing more: from s*d1 + t*d2 = 1, multiplying by n and
   dividing by d1*d2 gives n*s/d2 + n*t/d1. The remainders keep the
   numerator degrees below the denominator degrees. *)
let split_coprime (v : string) (n : Poly.t) (d1 : Poly.t) (d2 : Poly.t) : Poly.t * Poly.t =
  let g, s, t = Poly.ext_gcd v d1 d2 in
  if not (Poly.is_one g) then invalid_arg "split_coprime: denominators share a factor";
  let a = snd (Poly.quo_rem v (Poly.mul n t) d1) in
  let b = snd (Poly.quo_rem v (Poly.mul n s) d2) in
  (* n/(d1*d2) = a/d1 + b/d2, up to a polynomial that the caller has
     already divided out. *)
  (a, b)
 
(* n/q^k as a sum of c_j/q^j with deg c_j < deg q: repeatedly divide by
   q and read off the remainders, which is writing n in base q. *)
let expand_power (v : string) (n : Poly.t) (q : Poly.t) (k : int) : (Poly.t * int) list =
  let out = ref [] and cur = ref n in
  for j = k downto 1 do
    let quo, rem = Poly.quo_rem v !cur q in
    if not (Poly.is_zero rem) then out := (rem, j) :: !out;
    cur := quo
  done;
  List.rev !out

(* The full decomposition of a rational function in one variable:
   a polynomial part, then one term per irreducible factor power of the
   denominator. Irreducible over Q, which is why this needs the
   factoriser and not just a square-free split. *)
let apart (a : t) (v : string) : Poly.t * (Poly.t * Poly.t * int) list =
  if Poly.is_one a.den then (a.num, [])
  else begin
    let poly_part, rest = Poly.quo_rem v a.num a.den in
    if Poly.is_zero rest then (poly_part, [])
    else begin
      (* Factor the denominator over Q. *)
      let factors =
        match Poly.to_upoly v a.den with
        | None -> raise (Poly.Not_univariate "apart: denominator is not univariate")
        | Some (u, _) ->
          let _, fs = Factor.factor u in
          List.map (fun (g, m) -> (Poly.of_upoly v g, m)) fs
      in
      (* Peel the factor powers off one at a time with Bezout. Scale
         tracks the leading constant the factorisation dropped. *)
      let prod = List.fold_left (fun acc (g, m) -> Poly.mul acc (Poly.pow g m)) Poly.one factors in
      let scale =
        match (Poly.lead_term a.den, Poly.lead_term prod) with
        | Some (_, c1), Some (_, c2) -> Rational.div c1 c2
        | _ -> Rational.one
      in
      let numerator = ref (Poly.scale rest (Rational.inv scale)) in
      let remaining = ref prod in
      let out = ref [] in
      List.iter
        (fun (g, m) ->
          let gk = Poly.pow g m in
          let other = Poly.divide_exn !remaining gk in
          if Poly.is_one other then begin
            List.iter (fun (c, j) -> out := (c, g, j) :: !out) (expand_power v !numerator g m);
            numerator := Poly.zero;
            remaining := Poly.one
          end
          else begin
            let a1, a2 = split_coprime v !numerator gk other in
            List.iter (fun (c, j) -> out := (c, g, j) :: !out) (expand_power v a1 g m);
            numerator := a2;
            remaining := other
          end)
        factors;
      (poly_part, List.rev !out)
    end
  end

> apart 1/(x^2 - 1), x
-1/2*(1 + x)^(-1) + 1/2*(-1 + x)^(-1)
> apart (x^3)/((x - 1)^2), x
2 + x + (-1 + x)^(-2) + 3*(-1 + x)^(-1)
> apart 1/(x^3 + x), x
x^(-1) - x*(1 + x^2)^(-1)
> apart 1/((x - 1)^2*(x + 2)), x
-1/9*(-1 + x)^(-1) + 1/9*(2 + x)^(-1) + 1/3*(-1 + x)^(-2)

(* Factor, then read the roots off the factors. Degree one is exact,
   degree two is the quadratic formula, and anything irreducible of
   higher degree is reported as itself: RootOf(p, x) names a root
   without pretending there is a formula for it. *)
let solve (e : Expr.t) (v : string) : Expr.t list =
  let _, fs = factor_list e v in
  let c = new_ctx () in
  List.concat_map
    (fun (g, _) ->
      let p = to_poly c (Expr.simplify g) in
      match Poly.degree_in p v with
      | 1 ->
        let a = Poly.coeff_in p v 1 and b = Poly.coeff_in p v 0 in
        [ Expr.simplify (Expr.Mul [ of_poly c (Poly.neg b); Expr.Pow (of_poly c a, Expr.int (-1)) ]) ]
      | 2 ->
        let a = of_poly c (Poly.coeff_in p v 2) in
        let b = of_poly c (Poly.coeff_in p v 1) in
        let cc = of_poly c (Poly.coeff_in p v 0) in
        let disc =
          Expr.simplify
            (Expr.Add [ Expr.Pow (b, Expr.int 2); Expr.Mul [ Expr.int (-4); a; cc ] ])
        in
        let root = Expr.Pow (disc, Expr.rat 1 2) in
        let denom = Expr.Pow (Expr.Mul [ Expr.int 2; a ], Expr.int (-1)) in
        [ Expr.simplify (Expr.Mul [ Expr.Add [ Expr.Mul [ Expr.int (-1); b ]; root ]; denom ]);
          Expr.simplify
            (Expr.Mul
               [ Expr.Add [ Expr.Mul [ Expr.int (-1); b ]; Expr.Mul [ Expr.int (-1); root ] ]; denom ])
        ]
      | _ -> [ Expr.Fun ("RootOf", [ g; Expr.Sym v ]) ])
    fs

> solve x^2 - 5*x + 6, x
3, 2
> solve x^2 + x - 1, x
1/2*(-1 + 5^(1/2)), 1/2*(-1 - 5^(1/2))
> solve x^5 - x - 1, x
RootOf(-1 + x^5 - x, x)
> roots 6*x^2 - 5*x + 1, x
1/2, 1/3

(* Like to_poly, except that a negative integer power is a division
   rather than an opaque kernel. That one change is the difference
   between an expression tree and the field of fractions. *)
let rec to_ratfun (c : ctx) (e : Expr.t) : Ratfun.t =
  match e with
  | Expr.Num r -> Ratfun.of_poly (Poly.of_rational r)
  | Expr.Sym s ->
    ignore (kernel c (Expr.Sym s));
    Ratfun.of_poly (Poly.var s)
  | Expr.Add xs -> List.fold_left (fun acc x -> Ratfun.add acc (to_ratfun c x)) Ratfun.zero xs
  | Expr.Mul xs -> List.fold_left (fun acc x -> Ratfun.mul acc (to_ratfun c x)) Ratfun.one xs
  | Expr.Pow (b, Expr.Num n) when Rational.is_integer n -> (
    match Rational.to_int_opt n with
    | Some k -> Ratfun.pow (to_ratfun c b) k
    | None -> raise (Not_rational "exponent too large"))
  | Expr.Pow _ | Expr.Fun _ -> Ratfun.of_poly (Poly.var (kernel c e))
 
let of_ratfun (c : ctx) (r : Ratfun.t) : Expr.t =
  if Ratfun.is_polynomial r then of_poly c (Ratfun.num r)
  else
    Expr.simplify
      (Expr.Mul
         [ of_poly c (Ratfun.num r); Expr.Pow (of_poly c (Ratfun.den r), Expr.int (-1)) ])

(* Univariate over Q only: the multivariate case needs evaluation plus
   a multivariate Hensel lift, which is not in this part. *)
let factor_list (e : Expr.t) (v : string) : Rational.t * (Expr.t * int) list =
  let c = new_ctx () in
  let p = to_poly c (Expr.simplify e) in
  match Poly.to_upoly v p with
  | None -> raise (Not_factorable "factor: not univariate in that variable")
  | Some (u, scale) ->
    let cont, fs = Factor.factor u in
    let k = Rational.mul scale (Rational.of_bigint cont) in
    (k, List.map (fun (g, m) -> (of_poly c (Poly.of_upoly v g), m)) fs)
 
let factor (e : Expr.t) (v : string) : Expr.t =
  let k, fs = factor_list e v in
  let terms =
    List.map (fun (g, m) -> if m = 1 then g else Expr.Pow (g, Expr.int m)) fs
  in
  let terms = if Rational.is_one k then terms else Expr.Num k :: terms in
  match terms with [] -> Expr.Num k | [ t ] -> t | ts -> Expr.Mul ts

> factor x^2 - y^2, x
cannot factor: factor: not univariate in that variable

(* Randomized: plant a product of irreducibles and demand it back. *)
let test_factor_property () =
  Random.init 31337;
  let rec irreducible_of_degree d =
    let f =
      Upoly.of_bigint_list
        (List.init (d + 1) (fun i ->
           if i = d then Bigint.one else Bigint.of_int (Random.int 13 - 6)))
    in
    if Upoly.degree f = d && Factor.is_irreducible f then f else irreducible_of_degree d
  in
  for _ = 1 to 60 do
    let n = 2 + Random.int 2 in
    let planted = List.init n (fun _ -> irreducible_of_degree (1 + Random.int 3)) in
    let f = List.fold_left Upoly.mul Upoly.one planted in
    let c, fs = Factor.factor f in
    check "the factors multiply back to the input"
      (Upoly.equal
         (Upoly.scale (List.fold_left (fun acc (g, m) -> Upoly.mul acc (Upoly.pow g m)) Upoly.one fs) c)
         f);
    check "every factor is irreducible" (List.for_all (fun (g, _) -> Factor.is_irreducible g) fs);
    check "the total degree is preserved"
      (List.fold_left (fun acc (g, m) -> acc + (Upoly.degree g * m)) 0 fs = Upoly.degree f);
    (* Each planted factor has to divide the input, and so appear. *)
    check "every planted factor divides some returned factor"
      (List.for_all
         (fun g ->
           List.exists (fun (h, _) -> Upoly.divides (Upoly.primitive_part g) h <> None) fs)
         planted)
  done

(* The property: reassembling the decomposition has to give back the
   original fraction, exactly. That is checkable on random input, and
   it is what actually catches a wrong Bezout coefficient. *)
let test_apart_reassembles () =
  Random.init 4242;
  let random_linear () = add x (n (Random.int 9 - 4)) in
  let random_quadratic () = add (add (pow x 2) (mul (n (Random.int 3)) x)) (n (1 + Random.int 4)) in
  for _ = 1 to 120 do
    let den =
      List.fold_left
        (fun acc _ ->
          mul acc (if Random.bool () then random_linear () else random_quadratic ()))
        one
        (List.init (1 + Random.int 3) (fun i -> i))
    in
    let num =
      List.fold_left (fun acc k -> add acc (mul (n (Random.int 7 - 3)) (pow x k))) zero
        (List.init 4 (fun i -> i))
    in
    if (not (is_zero den)) && not (is_zero num) then begin
      let f = r num den in
      match Ratfun.apart f "x" with
      | exception Division_by_zero -> ()
      | poly_part, terms ->
        let rebuilt =
          List.fold_left
            (fun acc (c, g, j) -> Ratfun.add acc (r c (pow g j)))
            (Ratfun.of_poly poly_part) terms
        in
        check "partial fractions reassemble to the original" (Ratfun.equal rebuilt f);
        (* Every numerator must have degree below its denominator, or
           it was not fully decomposed. *)
        check "numerator degrees are below the denominator degrees"
          (List.for_all (fun (c, g, _) -> degree_in c "x" < degree_in g "x") terms)
    end
  done

  (* factor then expand has to be the identity *)
  List.iter
    (fun src ->
      check ("factor then expand is the identity: " ^ src)
        (Expr.compare_expr (Algebra.expand (Algebra.factor (p src) "x")) (Algebra.expand (p src)) = 0))
    [ "x^2 - 1"; "x^4 - 1"; "6*x^3 - 6"; "9*x^2 + 12*x + 4"; "x^5 - x"; "x^8 - 1";
      "x^4 - 8*x^2 + 12"; "2*x^3 + 3*x^2 - 2*x - 3" ];

multivariate factorization    factor x^2 - y^2 is refused. The route is
                              evaluation down to one variable, factor
                              there, then multivariate Hensel lifting
                              back up, with a leading-coefficient
                              correction that has no univariate analogue.
 
algebraic numbers             solve returns RootOf for degree 5 and
                              above, and sqrt for degree 2, but there is
                              no arithmetic on either. Q[x]/(p) would
                              give it.
 
no calculus at all            diff still refuses sin(x), there is no
                              series, no limit, no integral.
 
Part 5: derivatives with a function table, lazy power series,
        limits by Gruntz, and symbolic integration.