more nat bitwise ops

lib/Quox/NatExtra.idr

@@ -9,6 +9,11 @@ import Syntax.PreorderReasoning
 
 %default total
 
+infixl 8 `shiftL`, `shiftR`
+infixl 7 .&.
+infixl 6 `xor`
+infixl 5 .|.
+
 
 public export
 data LTE' : Nat -> Nat -> Type where
@@ -119,3 +124,86 @@ spread mask subj = go 1 (halfRec mask) subj 0 where
   go bit (HalfRecOdd  _ rec) subj res = case half subj of
     HalfOdd  subj => go (bit + bit) rec subj (res + bit)
     HalfEven subj => go (bit + bit) rec subj res
+
+
+
+public export
+data BitwiseRec : Nat -> Nat -> Type where
+  BwDone : BitwiseRec 0 0
+  Bw00   : (m, n : Nat) -> Lazy (BitwiseRec m n) ->
+           BitwiseRec (m + m) (n + n)
+  Bw01   : (m, n : Nat) -> Lazy (BitwiseRec m n) ->
+           BitwiseRec (m + m) (S (n + n))
+  Bw10   : (m, n : Nat) -> Lazy (BitwiseRec m n) ->
+           BitwiseRec (S (m + m)) (n + n)
+  Bw11   : (m, n : Nat) -> Lazy (BitwiseRec m n) ->
+           BitwiseRec (S (m + m)) (S (n + n))
+
+export
+bitwiseRec : (m, n : Nat) -> BitwiseRec m n
+bitwiseRec m n = go (halfRec m) (halfRec n) where
+  go : forall m, n. HalfRec m -> HalfRec n -> BitwiseRec m n
+  go HalfRecZ           HalfRecZ           = BwDone
+  go HalfRecZ           (HalfRecEven n nr) = Bw00 0 n $ go HalfRecZ nr
+  go HalfRecZ           (HalfRecOdd  n nr) = Bw01 0 n $ go HalfRecZ nr
+  go (HalfRecEven m mr) HalfRecZ           = Bw00 m 0 $ go mr HalfRecZ
+  go (HalfRecEven m mr) (HalfRecEven n nr) = Bw00 m n $ go mr nr
+  go (HalfRecEven m mr) (HalfRecOdd  n nr) = Bw01 m n $ go mr nr
+  go (HalfRecOdd  m mr) HalfRecZ           = Bw10 m 0 $ go mr HalfRecZ
+  go (HalfRecOdd  m mr) (HalfRecEven n nr) = Bw10 m n $ go mr nr
+  go (HalfRecOdd  m mr) (HalfRecOdd  n nr) = Bw11 m n $ go mr nr
+
+export
+bitwise : (Bool -> Bool -> Bool) -> Nat -> Nat -> Nat
+bitwise f m n = go 1 (bitwiseRec m n) 0 where
+  one : Bool -> Bool -> Nat -> Nat -> Nat
+  one p q bit res = if f p q then bit + res else res
+  go : forall m, n. Nat -> BitwiseRec m n -> Nat -> Nat
+  go bit BwDone         res = res
+  go bit (Bw00 m n rec) res = go (bit + bit) rec $ one False False bit res
+  go bit (Bw01 m n rec) res = go (bit + bit) rec $ one False True  bit res
+  go bit (Bw10 m n rec) res = go (bit + bit) rec $ one True  False bit res
+  go bit (Bw11 m n rec) res = go (bit + bit) rec $ one True  True  bit res
+
+export
+(.&.) : Nat -> Nat -> Nat
+(.&.) = bitwise $ \p, q => p && q
+
+private %foreign "scheme:blodwen-and"
+primAnd : Nat -> Nat -> Nat
+%transform "NatExtra.(.&.)" NatExtra.(.&.) m n = primAnd m n
+
+export
+(.|.) : Nat -> Nat -> Nat
+(.|.) = bitwise $ \p, q => p || q
+
+private %foreign "scheme:blodwen-or"
+primOr : Nat -> Nat -> Nat
+%transform "NatExtra.(.|.)" NatExtra.(.|.) m n = primOr m n
+
+export
+xor : Nat -> Nat -> Nat
+xor = bitwise (/=)
+
+private %foreign "scheme:blodwen-xor"
+primXor : Nat -> Nat -> Nat
+%transform "NatExtra.xor" NatExtra.xor m n = primXor m n
+
+
+export
+shiftL : Nat -> Nat -> Nat
+shiftL n 0     = n
+shiftL n (S i) = shiftL (n + n) i
+
+private %foreign "scheme:blodwen-shl"
+primShiftL : Nat -> Nat -> Nat
+%transform "NatExtra.shiftL" NatExtra.shiftL n i = primShiftL n i
+
+export
+shiftR : Nat -> Nat -> Nat
+shiftR n 0     = n
+shiftR n (S i) = shiftL (floorHalf n) i
+
+private %foreign "scheme:blodwen-shr"
+primShiftR : Nat -> Nat -> Nat
+%transform "NatExtra.shiftR" NatExtra.shiftR n i = primShiftR n i