(.module: [lux (#- and or not)])

(def: #export width Nat +64)

## [Values]
(do-template [<name> <op> <doc>]
  [(def: #export (<name> param subject)
     {#.doc <doc>}
     (All [s] (-> (I64 s) (I64 s) (I64 s)))
     (<op> param subject))]

  [and "lux i64 and" "Bitwise and."]
  [or  "lux i64 or"  "Bitwise or."]
  [xor "lux i64 xor" "Bitwise xor."]
  )

(do-template [<name> <op> <doc>]
  [(def: #export (<name> param subject)
     {#.doc <doc>}
     (All [s] (-> Nat (I64 s) (I64 s)))
     (<op> param subject))]

  [left-shift             "lux i64 left-shift"             "Bitwise left-shift."]
  [logical-right-shift    "lux i64 logical-right-shift"    "Unsigned bitwise logical-right-shift."]
  [arithmetic-right-shift "lux i64 arithmetic-right-shift" "Signed bitwise arithmetic-right-shift."]
  )

(def: (add-shift shift value)
  (-> Nat Nat Nat)
  (|> value (logical-right-shift shift) (n/+ value)))

(def: #export (count subject)
  {#.doc "Count the number of 1s in a bit-map."}
  (-> (I64 Any) Nat)
  (let [count' (n/- (|> subject (logical-right-shift +1) (..and +6148914691236517205) i64)
                    (i64 subject))]
    (|> count'
        (logical-right-shift +2) (..and +3689348814741910323) (n/+ (..and +3689348814741910323 count'))
        (add-shift +4) (..and +1085102592571150095)
        (add-shift +8)
        (add-shift +16)
        (add-shift +32)
        (..and +127))))

(def: #export not
  {#.doc "Bitwise negation."}
  (All [s] (-> (I64 s) (I64 s)))
  (xor (:coerce I64 -1)))

(def: (flag idx)
  (-> Nat I64)
  (|> +1 (:coerce I64) (left-shift idx)))

(def: #export (clear idx input)
  {#.doc "Clear bit at given index."}
  (All [s] (-> Nat (I64 s) (I64 s)))
  (|> idx flag ..not (..and input)))

(do-template [<name> <op> <doc>]
  [(def: #export (<name> idx input)
     {#.doc <doc>}
     (All [s] (-> Nat (I64 s) (I64 s)))
     (|> idx flag (<op> input)))]

  [set  ..or  "Set bit at given index."]
  [flip ..xor "Flip bit at given index."]
  )

(def: #export (set? idx input)
  (-> Nat (I64 Any) Bit)
  (|> input (:coerce I64) (..and (flag idx)) (n/= +0) .not))

(do-template [<name> <main> <comp>]
  [(def: #export (<name> distance input)
     (All [s] (-> Nat (I64 s) (I64 s)))
     (let [backwards-distance (n/- (n/% width distance) width)]
       (|> input
           (<comp> backwards-distance)
           (..or (<main> distance input)))))]

  [rotate-left  left-shift          logical-right-shift]
  [rotate-right logical-right-shift left-shift]
  )

(def: #export (region size offset)
  (-> Nat Nat I64)
  (|> +1 (:coerce I64) (left-shift size) dec (left-shift offset)))