Add non-zero restriction to gcd

libs/prelude/Prelude/Nat.idr

@@ -91,6 +91,12 @@ hyper (S pn)   a (S pb) = hyper pn a (hyper (S pn) a pb)
 -- Comparisons
 --------------------------------------------------------------------------------
 
+||| Proofs that `n` or `m` is not equal to Z
+data NotBothZero : (n, m : Nat) -> Type where
+  LeftIsNotZero  : NotBothZero (S n) m
+  RightIsNotZero : NotBothZero n     (S m)
+
+
 ||| Proofs that `n` is less than or equal to `m`
 ||| @ n the smaller number
 ||| @ m the larger number
@@ -386,16 +392,15 @@ log2 (S n) = log2NZ (S n) SIsNotZ
 --------------------------------------------------------------------------------
 -- GCD and LCM
 --------------------------------------------------------------------------------
-partial
-gcd : Nat -> Nat -> Nat
-gcd a Z = a
-gcd a b = assert_total (gcd b (a `modNat` b))
+gcd : (a: Nat) -> (b: Nat) -> .{auto ok: NotBothZero a b} -> Nat
+gcd a Z     = a
+gcd Z b     = b
+gcd a (S b) = assert_total $ gcd (S b) (modNatNZ a (S b) SIsNotZ)
 
-partial
 lcm : Nat -> Nat -> Nat
-lcm _ Z = Z
-lcm Z _ = Z
-lcm x y = divNat (x * y) (gcd x y)
+lcm _ Z     = Z
+lcm Z _     = Z
+lcm a (S b) = assert_total $ divNat (a * (S b)) (gcd a (S b))
 
 
 --------------------------------------------------------------------------------