-- |
-- Module:      Numeric.QuoteQuot
-- Copyright:   (c) 2020-2022 Andrew Lelechenko
-- Licence:     BSD3
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Generate routines for integer division, employing arithmetic
-- and bitwise operations only, which are __2.5x-3.5x faster__
-- than 'quot'. Divisors must be known in compile-time and be positive.
--

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE LexicalNegation #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}
{- HLINT ignore "Redundant lambda" -}

module Numeric.QuoteQuot
  (
  -- * Quasiquoters
    quoteQuot
  , quoteRem
  , quoteQuotRem
  , quoteDiv
  , quoteMod
  , quoteDivMod
  -- * AST
  , astQuot
  , AST(..)
  , interpretAST
  , quoteAST
  , assumeNonNegArg
  , MulHi(..)
  , Unsigned
  ) where

import Data.Bits (shiftL, shiftR, isSigned, FiniteBits, countTrailingZeros, finiteBitSize)
import Data.Int (Int8, Int16, Int32, Int64)
import Data.Word ( Word8, Word16, Word32, Word64 )
import GHC.Exts (dataToTag#, Int(..), Word(..), timesInt2#, timesWord2#)
import Language.Haskell.TH.Lift (Lift, Quote, Code)

-- | Quote integer division ('quot') by a compile-time known divisor, which
-- generates source code, employing arithmetic and bitwise operations only.
-- This is usually __2.5x-3.5x faster__ than using normal 'quot'.
--
-- > {-# LANGUAGE TemplateHaskell #-}
-- > {-# OPTIONS_GHC -ddump-splices -ddump-simpl -dsuppress-all #-}
-- > module Example where
-- > import Numeric.QuoteQuot
-- >
-- > -- Equivalent to (`quot` 10).
-- > quot10 :: Word -> Word
-- > quot10 = $$(quoteQuot 10)
--
-- >>> quot10 123
-- 12
--
-- Here @-ddump-splices@ demonstrates the chosen implementation
-- for division by 10:
--
-- > Splicing expression quoteQuot 10 ======>
-- > ((`shiftR` 3) . ((\ (W# w_a9N4) ->
-- >   let !(# hi_a9N5, _ #) = (timesWord2# w_a9N4) 14757395258967641293##
-- >   in W# hi_a9N5) . id))
--
-- And @-ddump-simpl@ demonstrates generated Core:
--
-- > quot10 = \ x_a5t2 ->
-- >   case x_a5t2 of { W# w_acHY ->
-- >   case timesWord2# w_acHY 14757395258967641293## of
-- >   { (# hi_acIg, ds_dcIs #) ->
-- >   W# (uncheckedShiftRL# hi_acIg 3#)
-- >   }
-- >   }
--
-- Benchmarks show that this implementation is __3.5x faster__
-- than @(`@'quot'@` 10)@.
--
quoteQuot :: forall a m. (MulHi a, Lift a, Quote m) => a -> Code m (a -> a)
quoteQuot :: forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
a -> Code m (a -> a)
quoteQuot a
d = AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST (a -> AST a
forall a. (Integral a, FiniteBits a) => a -> AST a
astQuot a
d)

-- | Similar to 'quoteQuot', but for 'rem'.
quoteRem :: (MulHi a, Lift a, Quote m) => a -> Code m (a -> a)
quoteRem :: forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
a -> Code m (a -> a)
quoteRem a
d = [|| (a, b) -> b
forall a b. (a, b) -> b
snd (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. $$(a -> Code m (a -> (a, a))
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
a -> Code m (a -> (a, a))
quoteQuotRem a
d) ||]

-- | Similar to 'quoteQuot', but for 'quotRem'.
quoteQuotRem :: (MulHi a, Lift a, Quote m) => a -> Code m (a -> (a, a))
quoteQuotRem :: forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
a -> Code m (a -> (a, a))
quoteQuotRem a
d = [|| \p
w -> let q :: a
q = $$(a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
a -> Code m (a -> a)
quoteQuot a
d) p
w in (a
q, p
w a -> a -> a
forall a. Num a => a -> a -> a
- a
d a -> a -> a
forall a. Num a => a -> a -> a
* a
q) ||]

-- | Similar to 'quoteQuot', but for 'div'.
--
-- @since 0.2.2.0
quoteDiv :: forall a m. (MulHi a, Lift a, MulHi (Unsigned a), Lift (Unsigned a), Quote m) => a -> Code m (a -> a)
quoteDiv :: forall a (m :: * -> *).
(MulHi a, Lift a, MulHi (Unsigned a), Lift (Unsigned a),
 Quote m) =>
a -> Code m (a -> a)
quoteDiv a
d
    | a -> Bool
forall a. Bits a => a -> Bool
isSigned a
d
    = [|| \(a'
i :: a') -> let w2i :: Unsigned a' -> a'
w2i = a -> b
Unsigned a' -> a'
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Unsigned a' -> a' in
                if a'
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0
                then Unsigned a' -> a'
w2i -(a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ $$Code m (Unsigned a -> Unsigned a)
go (a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral -(a'
i a -> a -> a
forall a. Num a => a -> a -> a
+ a
1)))
                else Unsigned a' -> a'
w2i (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
$ $$Code m (Unsigned a -> Unsigned a)
go (a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral a'
i) ||]
    | Bool
otherwise = a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
a -> Code m (a -> a)
quoteQuot a
d
  where
    go :: Code m (Unsigned a -> Unsigned a)
    go :: Code m (Unsigned a -> Unsigned a)
go = AST (Unsigned a) -> Code m (Unsigned a -> Unsigned a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST (Unsigned a -> AST (Unsigned a)
forall a. (Integral a, FiniteBits a) => a -> AST a
unsignedQuot (a -> Unsigned a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
d))

-- | Similar to 'quoteQuot', but for 'mod'.
--
-- @since 0.2.2.0
quoteMod :: (MulHi a, Lift a, MulHi (Unsigned a), Lift (Unsigned a), Quote m) => a -> Code m (a -> a)
quoteMod :: forall a (m :: * -> *).
(MulHi a, Lift a, MulHi (Unsigned a), Lift (Unsigned a),
 Quote m) =>
a -> Code m (a -> a)
quoteMod a
d = [|| (a, b) -> b
forall a b. (a, b) -> b
snd (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. $$(a -> Code m (a -> (a, a))
forall a (m :: * -> *).
(MulHi a, Lift a, MulHi (Unsigned a), Lift (Unsigned a),
 Quote m) =>
a -> Code m (a -> (a, a))
quoteDivMod a
d) ||]

-- | Similar to 'quoteQuot', but for 'divMod'.
--
-- @since 0.2.2.0
quoteDivMod :: (MulHi a, Lift a, MulHi (Unsigned a), Lift (Unsigned a), Quote m) => a -> Code m (a -> (a, a))
quoteDivMod :: forall a (m :: * -> *).
(MulHi a, Lift a, MulHi (Unsigned a), Lift (Unsigned a),
 Quote m) =>
a -> Code m (a -> (a, a))
quoteDivMod a
d = [|| \p
w -> let q :: a
q = $$(a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, MulHi (Unsigned a), Lift (Unsigned a),
 Quote m) =>
a -> Code m (a -> a)
quoteDiv a
d) p
w in (a
q, p
w a -> a -> a
forall a. Num a => a -> a -> a
- a
d a -> a -> a
forall a. Num a => a -> a -> a
* a
q) ||]

-- | Associate the corresponding unsigned type.
--
-- @since 0.2.2.0
type family Unsigned t
type instance Unsigned Int = Word
type instance Unsigned Word = Word
type instance Unsigned Int8 = Word8
type instance Unsigned Word8 = Word8
type instance Unsigned Int16 = Word16
type instance Unsigned Word16 = Word16
type instance Unsigned Int32 = Word32
type instance Unsigned Word32 = Word32
type instance Unsigned Int64 = Word64
type instance Unsigned Word64 = Word64

-- | Types allowing to multiply wide and return the high word of result.
class (Integral a, FiniteBits a) => MulHi a where
  mulHi :: a -> a -> a

instance MulHi Word8 where
  mulHi :: Word8 -> Word8 -> Word8
mulHi Word8
x Word8
y = Word16 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word8 -> Word16
widen Word8
x Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
* Word8 -> Word16
widen Word8
y) Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
`shiftR` Int
8)
    where widen :: Word8 -> Word16
widen = Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Word8 -> Word16

instance MulHi Word16 where
  mulHi :: Word16 -> Word16 -> Word16
mulHi Word16
x Word16
y = Word32 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word16 -> Word32
widen Word16
x Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* Word16 -> Word32
widen Word16
y) Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftR` Int
16)
    where widen :: Word16 -> Word32
widen = Word16 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Word16 -> Word32

instance MulHi Word32 where
  mulHi :: Word32 -> Word32 -> Word32
mulHi Word32
x Word32
y = Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word32 -> Word64
widen Word32
x Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word32 -> Word64
widen Word32
y) Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
32)
    where widen :: Word32 -> Word64
widen = Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Word32 -> Word64

-- | This instance is not efficient on 32-bit architecture.
instance MulHi Word64 where
  mulHi :: Word64 -> Word64 -> Word64
mulHi Word64
x Word64
y
    | Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
64
    = Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
x Word -> Word -> Word
forall a. MulHi a => a -> a -> a
`mulHi` Word64 -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
y :: Word)
    | Bool
otherwise
    = Word64 -> Word64 -> Word64
forall a. (Integral a, FiniteBits a) => a -> a -> a
defaultMulHi Word64
x Word64
y

instance MulHi Word where
  mulHi :: Word -> Word -> Word
mulHi (W# Word#
x) (W# Word#
y) = let !(# Word#
hi, Word#
_ #) = Word# -> Word# -> (# Word#, Word# #)
timesWord2# Word#
x Word#
y in Word# -> Word
W# Word#
hi

instance MulHi Int8 where
  mulHi :: Int8 -> Int8 -> Int8
mulHi Int8
x Int8
y = Int16 -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int8 -> Int16
widen Int8
x Int16 -> Int16 -> Int16
forall a. Num a => a -> a -> a
* Int8 -> Int16
widen Int8
y) Int16 -> Int -> Int16
forall a. Bits a => a -> Int -> a
`shiftR` Int
8)
    where widen :: Int8 -> Int16
widen = Int8 -> Int16
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Int8 -> Int16

instance MulHi Int16 where
  mulHi :: Int16 -> Int16 -> Int16
mulHi Int16
x Int16
y = Int32 -> Int16
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int16 -> Int32
widen Int16
x Int32 -> Int32 -> Int32
forall a. Num a => a -> a -> a
* Int16 -> Int32
widen Int16
y) Int32 -> Int -> Int32
forall a. Bits a => a -> Int -> a
`shiftR` Int
16)
    where widen :: Int16 -> Int32
widen = Int16 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Int16 -> Int32

instance MulHi Int32 where
  mulHi :: Int32 -> Int32 -> Int32
mulHi Int32
x Int32
y = Int64 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int32 -> Int64
widen Int32
x Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
* Int32 -> Int64
widen Int32
y) Int64 -> Int -> Int64
forall a. Bits a => a -> Int -> a
`shiftR` Int
32)
    where widen :: Int32 -> Int64
widen = Int32 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Int32 -> Int64

-- | This instance is not efficient on 32-bit architecture.
instance MulHi Int64 where
  mulHi :: Int64 -> Int64 -> Int64
mulHi Int64
x Int64
y
    | Int -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Int
0 :: Int) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
64
    = Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
x Int -> Int -> Int
forall a. MulHi a => a -> a -> a
`mulHi` Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
y :: Int)
    | Bool
otherwise
    = Int64 -> Int64 -> Int64
forall a. (Integral a, FiniteBits a) => a -> a -> a
defaultMulHi Int64
x Int64
y

instance MulHi Int where
  mulHi :: Int -> Int -> Int
mulHi (I# Int#
x) (I# Int#
y) = let !(# Int#
_, Int#
hi, Int#
_ #) = Int# -> Int# -> (# Int#, Int#, Int# #)
timesInt2# Int#
x Int#
y in Int# -> Int
I# Int#
hi

-- | An abstract syntax tree to represent
-- a function of one argument.
data AST a
  = Arg
  -- ^ Argument of the function
  | MulHi (AST a) a
  -- ^ Multiply wide and return the high word of result
  | MulLo (AST a) a
  -- ^ Multiply
  | Add (AST a) (AST a)
  -- ^ Add
  | Sub (AST a) (AST a)
  -- ^ Subtract
  | Shl (AST a) Int
  -- ^ Shift left
  | Shr (AST a) Int
  -- ^ Shift right with sign extension
  | CmpGE  (AST a) a
  -- ^ 1 if greater than or equal, 0 otherwise
  | CmpLT  (AST a) a
  -- ^ 1 if less than, 0 otherwise
  deriving (Int -> AST a -> ShowS
[AST a] -> ShowS
AST a -> String
(Int -> AST a -> ShowS)
-> (AST a -> String) -> ([AST a] -> ShowS) -> Show (AST a)
forall a. Show a => Int -> AST a -> ShowS
forall a. Show a => [AST a] -> ShowS
forall a. Show a => AST a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> AST a -> ShowS
showsPrec :: Int -> AST a -> ShowS
$cshow :: forall a. Show a => AST a -> String
show :: AST a -> String
$cshowList :: forall a. Show a => [AST a] -> ShowS
showList :: [AST a] -> ShowS
Show)

-- | Optimize 'AST', assuming that 'Arg' is non-negative.
assumeNonNegArg :: (Ord a, Num a) => AST a -> AST a
assumeNonNegArg :: forall a. (Ord a, Num a) => AST a -> AST a
assumeNonNegArg = \case
  Add AST a
x (CmpLT AST a
Arg a
n)
    | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 -> AST a
x
  Sub AST a
x (CmpLT AST a
Arg a
n)
    | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 -> AST a
x
  Add AST a
x (MulLo (CmpLT AST a
Arg a
n) a
_)
    | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0 -> AST a
x
  AST a
e -> AST a
e

-- | Reference (but slow) interpreter of 'AST'.
-- It is not meant to be used in production
-- and is provided primarily for testing purposes.
--
-- >>> interpretAST (astQuot (10 :: Data.Word.Word8)) 123
-- 12
--
interpretAST :: (Integral a, FiniteBits a) => AST a -> (a -> a)
interpretAST :: forall a. (Integral a, FiniteBits a) => AST a -> a -> a
interpretAST AST a
ast a
n = AST a -> a
go AST a
ast
  where
    go :: AST a -> a
go = \case
      AST a
Arg       -> a
n
      MulHi AST a
x a
k -> a -> a -> a
forall a. (Integral a, FiniteBits a) => a -> a -> a
defaultMulHi (AST a -> a
go AST a
x) a
k
      MulLo AST a
x a
k -> AST a -> a
go AST a
x a -> a -> a
forall a. Num a => a -> a -> a
* a
k
      Add AST a
x AST a
y   -> AST a -> a
go AST a
x a -> a -> a
forall a. Num a => a -> a -> a
+ AST a -> a
go AST a
y
      Sub AST a
x AST a
y   -> AST a -> a
go AST a
x a -> a -> a
forall a. Num a => a -> a -> a
- AST a -> a
go AST a
y
      Shl AST a
x Int
k   -> AST a -> a
go AST a
x a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
k
      Shr AST a
x Int
k   -> AST a -> a
go AST a
x a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
k
      CmpGE AST a
x a
k -> if AST a -> a
go AST a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
k then a
1 else a
0
      CmpLT AST a
x a
k -> if AST a -> a
go AST a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<  a
k then a
1 else a
0

defaultMulHi :: (Integral a, FiniteBits a) => a -> a -> a
defaultMulHi :: forall a. (Integral a, FiniteBits a) => a -> a -> a
defaultMulHi a
x a
y =
    Integer -> a
forall a. Num a => Integer -> a
fromInteger (Integer -> a) -> Integer -> a
forall a b. (a -> b) -> a -> b
$ (a -> Integer
forall a. Integral a => a -> Integer
toInteger a
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* a -> Integer
forall a. Integral a => a -> Integer
toInteger a
y) Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftR` a -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize a
x

iprop :: Integral a => (a -> Bool) -> a -> a
iprop :: forall a. Integral a => (a -> Bool) -> a -> a
iprop a -> Bool
p = \a
i -> Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int# -> Int
I# (Bool -> Int#
forall a. DataToTag a => a -> Int#
dataToTag# (a -> Bool
p a
i)))
{-# INLINE iprop #-}


-- | Embed 'AST' into Haskell expression.
quoteAST :: (MulHi a, Lift a, Quote m) => AST a -> Code m (a -> a)
quoteAST :: forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST = \case
  AST a
Arg            -> [|| a -> a
forall a. a -> a
id ||]
  Shr AST a
x Int
k        -> [|| (a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
k) (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
x) ||]
  Shl AST a
x Int
k        -> [|| (a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
k) (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
x) ||]
  MulHi AST a
x a
k      -> [|| (a -> a -> a
forall a. MulHi a => a -> a -> a
`mulHi` a
k) (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
x) ||]
  MulLo AST a
x a
k      -> [|| (a -> a -> a
forall a. Num a => a -> a -> a
* a
k) (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
x) ||]
  Add AST a
x AST a
y        -> [|| \p
w -> $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
x) p
w a -> a -> a
forall a. Num a => a -> a -> a
+ $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
y) p
w ||]
  Sub AST a
x AST a
y        -> [|| \p
w -> $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
x) p
w a -> a -> a
forall a. Num a => a -> a -> a
- $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
y) p
w ||]
  CmpGE AST a
x a
k      -> [|| (a -> Bool) -> a -> a
forall a. Integral a => (a -> Bool) -> a -> a
iprop (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
k) (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
x) ||]
  CmpLT AST a
x a
k      -> [|| (a -> Bool) -> a -> a
forall a. Integral a => (a -> Bool) -> a -> a
iprop (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
k) (b -> c) -> (a -> b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. $$(AST a -> Code m (a -> a)
forall a (m :: * -> *).
(MulHi a, Lift a, Quote m) =>
AST a -> Code m (a -> a)
quoteAST AST a
x) ||]

-- | 'astQuot' @d@ constructs an 'AST' representing
-- a function, equivalent to 'quot' @a@ for positive @a@,
-- but avoiding division instructions.
--
-- >>> astQuot (10 :: Data.Word.Word8)
-- Shr (MulHi Arg 205) 3
--
-- And indeed to divide 'Data.Word.Word8' by 10
-- one can multiply it by 205, take the high byte and
-- shift it right by 3. Somewhat counterintuitively,
-- this sequence of operations is faster than a single
-- division on most modern achitectures.
--
-- 'astQuot' function is polymorphic and supports both signed
-- and unsigned operands of arbitrary finite bitness.
-- Implementation is based on
-- Ch. 10 of Hacker's Delight by Henry S. Warren, 2012.
--
astQuot :: (Integral a, FiniteBits a) => a -> AST a
astQuot :: forall a. (Integral a, FiniteBits a) => a -> AST a
astQuot a
k
  | a -> Bool
forall a. Bits a => a -> Bool
isSigned a
k = a -> AST a
forall a. (Integral a, FiniteBits a) => a -> AST a
signedQuot a
k
  | Bool
otherwise  = a -> AST a
forall a. (Integral a, FiniteBits a) => a -> AST a
unsignedQuot a
k

unsignedQuot :: (Integral a, FiniteBits a) => a -> AST a
unsignedQuot :: forall a. (Integral a, FiniteBits a) => a -> AST a
unsignedQuot a
k'
  | a -> Bool
forall a. Bits a => a -> Bool
isSigned a
k
  = String -> AST a
forall a. HasCallStack => String -> a
error String
"unsignedQuot works for unsigned types only"
  | a
k' a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0
  = String -> AST a
forall a. HasCallStack => String -> a
error String
"divisor must be positive"
  | a
k' a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1
  = AST a
forall a. AST a
Arg
  | a
k a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1
  = AST a -> Int -> AST a
forall a. AST a -> Int -> AST a
shr AST a
forall a. AST a
Arg Int
kZeros
  | a
k' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
1 a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` (Int
fbs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  = AST a -> a -> AST a
forall a. AST a -> a -> AST a
CmpGE AST a
forall a. AST a
Arg a
k'
  -- Hacker's Delight, 10-8, Listing 1
  | a
k a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
1 a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
shft
  = AST a -> Int -> AST a
forall a. AST a -> Int -> AST a
shr (AST a -> a -> AST a
forall a. AST a -> a -> AST a
MulHi AST a
forall a. AST a
Arg a
magic) Int
shft'
  -- Hacker's Delight, 10-8, Listing 3
  | Bool
otherwise
  = AST a -> Int -> AST a
forall a. AST a -> Int -> AST a
shr (AST a -> AST a -> AST a
forall a. AST a -> AST a -> AST a
Add (AST a -> Int -> AST a
forall a. AST a -> Int -> AST a
shr (AST a -> AST a -> AST a
forall a. AST a -> AST a -> AST a
Sub AST a
forall a. AST a
Arg (AST a -> a -> AST a
forall a. AST a -> a -> AST a
MulHi AST a
forall a. AST a
Arg a
magic)) Int
1) (AST a -> a -> AST a
forall a. AST a -> a -> AST a
MulHi AST a
forall a. AST a
Arg a
magic)) (Int
shft' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    fbs :: Int
fbs = a -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize a
k'
    kZeros :: Int
kZeros = a -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros a
k'
    k :: a
k = a
k' a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
kZeros
    r0 :: a
r0 = Integer -> a
forall a. Num a => Integer -> a
fromInteger ((Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` Int
fbs) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` a -> Integer
forall a. Integral a => a -> Integer
toInteger a
k)
    shft :: Int
shft = a -> Int -> Int
go a
r0 Int
0
    shft' :: Int
shft' = Int
shft Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kZeros
    magic :: a
magic = Integer -> a
forall a. Num a => Integer -> a
fromInteger ((Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` (Int
fbs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
shft)) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` a -> Integer
forall a. Integral a => a -> Integer
toInteger a
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)

    go :: a -> Int -> Int
go a
r Int
s
      | (a
k a -> a -> a
forall a. Num a => a -> a -> a
- a
r) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
1 a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
s = Int
s
      | Bool
otherwise = a -> Int -> Int
go (a
r a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
1 a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
k) (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

signedQuot :: (Integral a, FiniteBits a) => a -> AST a
signedQuot :: forall a. (Integral a, FiniteBits a) => a -> AST a
signedQuot a
k'
  | Bool -> Bool
not (a -> Bool
forall a. Bits a => a -> Bool
isSigned a
k)
  = String -> AST a
forall a. HasCallStack => String -> a
error String
"signedQuot works for signed types only"
  | a
k' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0
  = String -> AST a
forall a. HasCallStack => String -> a
error String
"divisor must be positive"
  | a
k' a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1
  = AST a
forall a. AST a
Arg
  -- Hacker's Delight, 10-1, Listing 2
  | a
k a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1
  = AST a -> Int -> AST a
forall a. AST a -> Int -> AST a
shr (AST a -> AST a -> AST a
forall a. AST a -> AST a -> AST a
Add AST a
forall a. AST a
Arg (AST a -> a -> AST a
forall a. AST a -> a -> AST a
MulLo (AST a -> a -> AST a
forall a. AST a -> a -> AST a
CmpLT AST a
forall a. AST a
Arg a
0) (a
k' a -> a -> a
forall a. Num a => a -> a -> a
- a
1))) Int
kZeros
  | a
k' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
1 a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` (Int
fbs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2)
  = AST a -> AST a -> AST a
forall a. AST a -> AST a -> AST a
Sub (AST a -> a -> AST a
forall a. AST a -> a -> AST a
CmpGE AST a
forall a. AST a
Arg a
k') (AST a -> a -> AST a
forall a. AST a -> a -> AST a
CmpLT AST a
forall a. AST a
Arg (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
k'))
  -- Hacker's Delight, 10-3, Listing 2
  | a
magic a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0
  = AST a -> AST a -> AST a
forall a. AST a -> AST a -> AST a
Add (AST a -> Int -> AST a
forall a. AST a -> Int -> AST a
shr (AST a -> a -> AST a
forall a. AST a -> a -> AST a
MulHi AST a
forall a. AST a
Arg a
magic) Int
shft') (AST a -> a -> AST a
forall a. AST a -> a -> AST a
CmpLT AST a
forall a. AST a
Arg a
0)
  -- Hacker's Delight, 10-3, Listing 3
  | Bool
otherwise
  = AST a -> AST a -> AST a
forall a. AST a -> AST a -> AST a
Add (AST a -> Int -> AST a
forall a. AST a -> Int -> AST a
shr (AST a -> AST a -> AST a
forall a. AST a -> AST a -> AST a
Add AST a
forall a. AST a
Arg (AST a -> a -> AST a
forall a. AST a -> a -> AST a
MulHi AST a
forall a. AST a
Arg a
magic)) Int
shft') (AST a -> a -> AST a
forall a. AST a -> a -> AST a
CmpLT AST a
forall a. AST a
Arg a
0)
  where
    fbs :: Int
fbs = a -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize a
k'
    kZeros :: Int
kZeros = a -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros a
k'
    k :: a
k = a
k' a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
kZeros
    r0 :: a
r0 = Integer -> a
forall a. Num a => Integer -> a
fromInteger ((Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` Int
fbs) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` a -> Integer
forall a. Integral a => a -> Integer
toInteger a
k)
    shft :: Int
shft = a -> Int -> Int
go a
r0 Int
0
    shft' :: Int
shft' = Int
shft Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
kZeros
    magic :: a
magic = Integer -> a
forall a. Num a => Integer -> a
fromInteger ((Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` (Int
fbs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
shft)) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`quot` a -> Integer
forall a. Integral a => a -> Integer
toInteger a
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)

    go :: a -> Int -> Int
go a
r Int
s
      | (a
k a -> a -> a
forall a. Num a => a -> a -> a
- a
r) a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
1 a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) = Int
s
      | Bool
otherwise = a -> Int -> Int
go (a
r a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
1 a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
k) (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

shr :: AST a -> Int -> AST a
shr :: forall a. AST a -> Int -> AST a
shr AST a
x Int
0 = AST a
x
shr AST a
x Int
k = AST a -> Int -> AST a
forall a. AST a -> Int -> AST a
Shr AST a
x Int
k