{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}

-- | Utilities for working with 'Word' and 'Word#'.
module Agda.Utils.Word
  ( -- * Bitwise operations
    uncheckedBitWord#
  , uncheckedClearBitWord#
  , uncheckedSetBitWord#
  , uncheckedTestBitWord#
  , highestBitWord#
  , lowestBitWord#
  , disjointWord#
  , andNot#
  , uncheckedWordOnes#
  -- * Folds
  --
  -- $wordFolds
  , wordFoldrBits#
  , wordFoldlBits#
  , wordFoldrBitsOffset#
  , wordFoldlBitsOffset#
  -- ** Strict folds
  , wordFoldrBitsStrict#
  , wordFoldlBitsStrict#
  , wordFoldrBitsOffsetStrict#
  , wordFoldlBitsOffsetStrict#
  ) where

-- We need the machines word size for some bitwise operations.
#include "MachDeps.h"

import GHC.Base

--------------------------------------------------------------------------------
-- Bitwise operations

-- | Create a 'Word#' with the @i@th bit set.
--
-- The result of 'uncheckedBitWord#' is undefined when @i@
-- is not in the range @[0..WORD_SIZE_IN_BITS]@.
uncheckedBitWord# :: Int# -> Word#
uncheckedBitWord# :: Int# -> Word#
uncheckedBitWord# Int#
i = Word#
1## Word# -> Int# -> Word#
`uncheckedShiftL#` Int#
i
{-# INLINE uncheckedBitWord# #-}

-- | Clear a single bit in an unboxed word.
--
-- The result of 'uncheckedClearBitWord#' is undefined when @i@
-- is not in the range @[0..WORD_SIZE_IN_BITS]@.
uncheckedClearBitWord# :: Word# -> Int# -> Word#
uncheckedClearBitWord# :: Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i =
  -- This is somewhat suboptimal. On x86, the GHC native code generator will
  -- produce the following code with -O2:
  --
  -- > movl $1,%eax
  -- > movq %rsi,%rcx
  -- > shlq %cl,%rax
  -- > notq %rax
  -- > movq %r14,%rbx
  -- > andq %rax,%rbx
  -- > jmp *(%rbp)
  --
  -- Ideally, this should be a single @btr@ instruction (or @bic@ on arm), but
  -- GHC doesn't provide any builtins for this so this is the best we can do.
  Word#
w Word# -> Word# -> Word#
`and#` (Word# -> Word#
not# (Word#
1## Word# -> Int# -> Word#
`uncheckedShiftL#` Int#
i))
{-# INLINE uncheckedClearBitWord# #-}

-- | Set a single bit in an unboxed word.
--
-- The result of 'uncheckedSetBitWord#' is undefined when @i@
-- is not in the range @[0..WORD_SIZE_IN_BITS]@.
uncheckedSetBitWord# :: Word# -> Int# -> Word#
uncheckedSetBitWord# :: Word# -> Int# -> Word#
uncheckedSetBitWord# Word#
w Int#
i =
  Word#
w Word# -> Word# -> Word#
`or#` (Word#
1## Word# -> Int# -> Word#
`uncheckedShiftL#` Int#
i)
{-# INLINE uncheckedSetBitWord# #-}

-- | Test a single bit in an unboxed word.
--
-- The result of 'uncheckedTestBitWord#' is undefined when @i@
-- is not in the range @[0..WORD_SIZE_IN_BITS]@.
uncheckedTestBitWord# :: Word# -> Int# -> Int#
uncheckedTestBitWord# :: Word# -> Int# -> Int#
uncheckedTestBitWord# Word#
w Int#
i =
  (Word#
w Word# -> Word# -> Word#
`and#` (Word#
1## Word# -> Int# -> Word#
`uncheckedShiftL#` Int#
i)) Word# -> Word# -> Int#
`neWord#` Word#
0##
{-# INLINE uncheckedTestBitWord# #-}

-- | Get the bit-index of the highest set bit.
highestBitWord# :: Word# -> Word#
highestBitWord# :: Word# -> Word#
highestBitWord# Word#
w =
  -- Another bit of suboptimal code resulting in a lack of primops.
  -- If we compile with -O2 using the native code generator on an x86 machine,
  -- we get:
  --
  -- > bsr %r14,%rax
  -- > movl $127,%ebx
  -- > cmovne %rax,%rbx
  -- > xorq $63,%rbx
  -- > xorq $63,%rbx
  -- > jmp *(%rbp)
  --
  -- Note the double @xor@! This could be a single @bsr@ instruction.
  -- If we compile with @-mbmi2@ things are a bit better, but still wastes
  -- an instruction. However, this is the best we can do.
  Word# -> Word#
clz# Word#
w Word# -> Word# -> Word#
`xor#` (WORD_SIZE_IN_BITS## `minusWord#` 1##)
{-# INLINE highestBitWord# #-}

-- | Get the bit-index of the lowest set bit.
lowestBitWord# :: Word# -> Word#
lowestBitWord# :: Word# -> Word#
lowestBitWord# Word#
w = Word# -> Word#
ctz# Word#
w
{-# INLINE lowestBitWord# #-}

-- | Take the bitwise AND of the first argument with the complement of the second.
andNot# :: Word# -> Word# -> Word#
andNot# :: Word# -> Word# -> Word#
andNot# Word#
w1 Word#
w2 =
  -- This should be a single @andn@ on x86 or @bic@ on aarch64,
  -- but the native code generator doesn't know about this.
  Word#
w1 Word# -> Word# -> Word#
`and#` Word# -> Word#
not# Word#
w2
{-# INLINE andNot# #-}

-- | Are the bit representations of two words disjoint?
disjointWord# :: Word# -> Word# -> Int#
disjointWord# :: Word# -> Word# -> Int#
disjointWord# Word#
w1 Word#
w2 = ((Word#
w1 Word# -> Word# -> Word#
`and#` Word#
w2) Word# -> Word# -> Int#
`eqWord#` Word#
0##)
{-# INLINE disjointWord# #-}

-- | Create a 'Word#' where the first @n@ bits are 1.
--
-- The result of 'uncheckedWordOnes#' is undefined when @i@
-- is not in the range @[0..WORD_SIZE_IN_BITS]@.
uncheckedWordOnes# :: Int# -> Word#
uncheckedWordOnes# :: Int# -> Word#
uncheckedWordOnes# Int#
n = Word# -> Word#
not# Word#
0## Word# -> Int# -> Word#
`uncheckedShiftRL#` (WORD_SIZE_IN_BITS# -# n)
{-# INLINE uncheckedWordOnes# #-}

--------------------------------------------------------------------------------
-- Folds

-- $wordFolds
-- As usual, there is an ambiguity in left/right folds for folding over the bits of a
-- 'Word#'. We opt to use the convention where we treat the 0th bit as the "head" of
-- a 'Word#', so a right fold like @wordFoldrBits# f x 0b1011@ would give @f 0 (f 1 (f 3 x))@.
--
-- We also provide hand-rolled versions of folds that offset the bit index by a static value:
-- this pattern comes up quite often when folding over 'ByteArray#' and the like. This can
-- save some allocations when compared to @wordFoldrBits# (\i acc -> f (i + offset) acc) a@.

-- | @wordFoldrBits# f a w@ performs a lazy right fold over the bits of @w@.
wordFoldrBits# :: (Int -> a -> a) -> a -> Word# -> a
wordFoldrBits# :: forall a. (Int -> a -> a) -> a -> Word# -> a
wordFoldrBits# Int -> a -> a
f a
a Word#
w
  | Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) = a
a
  | Bool
otherwise =
    let i :: Int#
i = Word# -> Int#
word2Int# (Word# -> Word#
lowestBitWord# Word#
w)
    in Int -> a -> a
f (Int# -> Int
I# Int#
i) ((Int -> a -> a) -> a -> Word# -> a
forall a. (Int -> a -> a) -> a -> Word# -> a
wordFoldrBits# Int -> a -> a
f a
a (Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i))

-- | @wordFoldrBitsOffset# offset f a w@ performs a right fold over the bits of @w@,
-- with every bit index passed to @f@ offset by @offset@.
wordFoldrBitsOffset# :: Int# -> (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsOffset# :: forall a. Int# -> (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsOffset# Int#
offset Int -> a -> a
f a
a Word#
w
  | Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) = a
a
  | Bool
otherwise =
    let i :: Int#
i = Word# -> Int#
word2Int# (Word# -> Word#
lowestBitWord# Word#
w)
    in Int -> a -> a
f (Int# -> Int
I# (Int#
i Int# -> Int# -> Int#
+# Int#
offset)) (Int# -> (Int -> a -> a) -> a -> Word# -> a
forall a. Int# -> (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsOffset# Int#
offset Int -> a -> a
f a
a (Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i))

-- | @wordFoldlBits# f a w@ performs a lazy left fold over the bits of @w@.
wordFoldlBits# :: (a -> Int -> a) -> a -> Word# -> a
wordFoldlBits# :: forall a. (a -> Int -> a) -> a -> Word# -> a
wordFoldlBits# a -> Int -> a
f a
a Word#
w
  | Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) = a
a
  | Bool
otherwise =
      let i :: Int#
i = Word# -> Int#
word2Int# (Word# -> Word#
highestBitWord# Word#
w)
      in a -> Int -> a
f ((a -> Int -> a) -> a -> Word# -> a
forall a. (a -> Int -> a) -> a -> Word# -> a
wordFoldlBits# a -> Int -> a
f a
a (Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i)) (Int# -> Int
I# Int#
i)

-- | @wordFoldlBitsOffset# offset f a w@ performs a lazy left fold over the bits of @w@,
-- with every bit index passed to @f@ offset by @offset@.
wordFoldlBitsOffset# :: Int# -> (a -> Int -> a) -> a -> Word# -> a
wordFoldlBitsOffset# :: forall a. Int# -> (a -> Int -> a) -> a -> Word# -> a
wordFoldlBitsOffset# Int#
offset a -> Int -> a
f a
a Word#
w
  | Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) = a
a
  | Bool
otherwise =
      let i :: Int#
i = Word# -> Int#
word2Int# (Word# -> Word#
highestBitWord# Word#
w)
      in a -> Int -> a
f (Int# -> (a -> Int -> a) -> a -> Word# -> a
forall a. Int# -> (a -> Int -> a) -> a -> Word# -> a
wordFoldlBitsOffset# Int#
offset a -> Int -> a
f a
a (Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i)) (Int# -> Int
I# (Int#
i Int# -> Int# -> Int#
+# Int#
offset))

-- | @wordFoldrBitsStrict# f a w@ performs a strict right fold over the bits of @w@.
wordFoldrBitsStrict# :: (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsStrict# :: forall a. (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsStrict# Int -> a -> a
f !a
a Word#
w
  | Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) = a
a
  | Bool
otherwise =
    let i :: Int#
i = Word# -> Int#
word2Int# (Word# -> Word#
highestBitWord# Word#
w)
    in (Int -> a -> a) -> a -> Word# -> a
forall a. (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsStrict# Int -> a -> a
f (Int -> a -> a
f (Int# -> Int
I# Int#
i) a
a) (Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i)

-- | @wordFoldrBitsOffsetStrict# offset f a w@ performs a strict right fold over the bits of @w@,
-- with every bit index passed to @f@ offset by @offset@.
wordFoldrBitsOffsetStrict# :: Int# -> (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsOffsetStrict# :: forall a. Int# -> (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsOffsetStrict# Int#
offset Int -> a -> a
f !a
a Word#
w
  | Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) = a
a
  | Bool
otherwise =
    let i :: Int#
i = Word# -> Int#
word2Int# (Word# -> Word#
highestBitWord# Word#
w)
    in Int# -> (Int -> a -> a) -> a -> Word# -> a
forall a. Int# -> (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsOffsetStrict# Int#
offset Int -> a -> a
f (Int -> a -> a
f (Int# -> Int
I# (Int#
i Int# -> Int# -> Int#
+# Int#
offset)) a
a) (Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i)

-- | @wordFoldlBitsStrict# f a w@ performs a strict left fold over the bits of @w@.
wordFoldlBitsStrict# :: (a -> Int -> a) -> a -> Word# -> a
wordFoldlBitsStrict# :: forall a. (a -> Int -> a) -> a -> Word# -> a
wordFoldlBitsStrict# a -> Int -> a
f !a
a Word#
w
  | Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) = a
a
  | Bool
otherwise =
    let i :: Int#
i = Word# -> Int#
word2Int# (Word# -> Word#
lowestBitWord# Word#
w)
    in (a -> Int -> a) -> a -> Word# -> a
forall a. (a -> Int -> a) -> a -> Word# -> a
wordFoldlBitsStrict# a -> Int -> a
f (a -> Int -> a
f a
a (Int# -> Int
I# Int#
i)) (Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i)

-- | @wordFoldrBitsOffsetStrict# offset f a w@ performs a strict right fold over the bits of @w@,
-- with every bit index passed to @f@ offset by @offset@.
wordFoldlBitsOffsetStrict# :: Int# -> (a -> Int -> a) -> a -> Word# -> a
wordFoldlBitsOffsetStrict# :: forall a. Int# -> (a -> Int -> a) -> a -> Word# -> a
wordFoldlBitsOffsetStrict# Int#
offset a -> Int -> a
f !a
a Word#
w
  | Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) = a
a
  | Bool
otherwise =
    let i :: Int#
i = Word# -> Int#
word2Int# (Word# -> Word#
lowestBitWord# Word#
w)
    in Int# -> (a -> Int -> a) -> a -> Word# -> a
forall a. Int# -> (a -> Int -> a) -> a -> Word# -> a
wordFoldlBitsOffsetStrict# Int#
offset a -> Int -> a
f (a -> Int -> a
f a
a (Int# -> Int
I# (Int#
i Int# -> Int# -> Int#
+# Int#
offset))) (Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i)