Something went wrong. Try again.
dependent-types proof-assistant
Something went wrong. Try again.
8.6 kB · 278 lines
Haskell
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279
-- | Persistent array using SmallArray# nodes and a tail cache, in the style-- of Clojure persistent arrays.---- The branching factor W is chosen with CPP: compile with -DBRANCHING=4, 8 or-- 16. The default is 8.module Utils.TreeArray ( Array (..) , empty , index , modify , size , snoc , update , same ) where
#ifndef BRANCHING#define BRANCHING 8#endif
import GHC.Exts hiding (toList)import Data.Bitsimport Data.Kind (Type)
---------------------------------------------------------------------------------- Branching-dependent definitions
-- | log2 of the branching factor.bits :: Int
-- | Shift of the root node for a tree of the given (positive) number of-- chunks: the bit length of the largest chunk index (m - 1), rounded up to a-- multiple of "bits", minus "bits". A 1-chunk tree has shift (-bits): its-- root is the chunk itself.rootShift :: Int -> Int
-- | Every "bits"-th bit set: m is a power of W iff it is a power of two and-- has one of these bits set (this also rules out m == 0).powMask :: Int
#if BRANCHING == 4bits = 2rootShift m = ((65 - countLeadingZeros (m - 1)) .&. (-2)) - 2powMask = 0x5555555555555555#elif BRANCHING == 8bits = 3-- x `quot` 3 == (x * 43) >> 7 for 0 <= x < 128, and here 2 <= x <= 65;-- q * 3 == q + q << 1. GHC's native code generator would emit idiv and imul.rootShift m = let q = unsafeShiftR ((66 - countLeadingZeros (m - 1)) * 43) 7 in q + unsafeShiftL q 1 - 3powMask = 0x1249249249249249#elif BRANCHING == 16bits = 4rootShift m = ((67 - countLeadingZeros (m - 1)) .&. (-4)) - 4powMask = 0x1111111111111111#else#error "BRANCHING must be 4, 8 or 16"#endif
{-# inline bits #-}{-# inline rootShift #-}{-# inline powMask #-}
{-# inline width #-}width :: Intwidth = unsafeShiftL 1 bits
{-# inline mask #-}mask :: Intmask = width - 1
---------------------------------------------------------------------------------- Nodes
type U = (Any :: UnliftedType)type Node = SmallArray# U
type role Array representational-- | Chunk tree, tail, size.data Array (a :: Type) = Array U Node Int
{-# inline toU #-}toU :: a -> UtoU x = unsafeCoerce# x -- x is forced by Strict
{-# inline fromU #-}fromU :: U -> afromU x = unsafeCoerce# x
{-# inline nodeU #-}nodeU :: Node -> UnodeU n = unsafeCoerce# n
{-# inline uNode #-}uNode :: U -> NodeuNode t = unsafeCoerce# t
-- | Filler for unused slots.{-# inline hole #-}hole :: (# #) -> Uhole _ = toU ()
{-# inline get #-}get :: Node -> Int -> Uget n (I# i) = case indexSmallArray# n i of (# x #) -> x
{-# inline set #-}-- | Copy of the node with one child replaced.set :: Node -> Int -> U -> Nodeset n (I# i) v = case width of I# w -> runRW# \s -> case thawSmallArray# n 0# w s of (# s, m #) -> case writeSmallArray# m i v s of s -> case unsafeFreezeSmallArray# m s of (# _, n' #) -> n'
{-# inline node1 #-}-- | Node with the given first child and holes elsewhere.node1 :: U -> Nodenode1 a = case width of I# w -> runRW# \s -> case newSmallArray# w (hole (# #)) s of (# s, m #) -> case writeSmallArray# m 0# a s of s -> case unsafeFreezeSmallArray# m s of (# _, n #) -> n
{-# inline node2 #-}-- | Node with the given first two children and holes elsewhere.node2 :: U -> U -> Nodenode2 a b = case width of I# w -> runRW# \s -> case newSmallArray# w (hole (# #)) s of (# s, m #) -> case writeSmallArray# m 0# a s of s -> case writeSmallArray# m 1# b s of s -> case unsafeFreezeSmallArray# m s of (# _, n #) -> n
---------------------------------------------------------------------------------- Operations
empty :: Array aempty = Array (hole (# #)) (node1 (hole (# #))) 0
-- | 0 <= x < s, as a single unsigned comparison.inBounds :: Int -> Int -> BoolinBounds x s = (fromIntegral x :: Word) < fromIntegral s
{-# inline digit #-}digit :: Int -> Int -> Intdigit x shift = unsafeShiftR x shift .&. mask
{-# inline size #-}size :: Array a -> Intsize (Array _ _ s) = s
{-# inline same #-}same :: Array a -> Array a -> Boolsame (Array t tl _) (Array t' tl' _) = isTrue# (reallyUnsafePtrEquality# t t') && isTrue# (reallyUnsafePtrEquality# tl tl')
{-# inline treeSize #-}-- | Number of elements stored in the tree.treeSize :: Int -> InttreeSize s = s .&. complement mask
-- | Chunk number c of a tree with root at "shift".chunk :: U -> Int -> Int -> Nodechunk t shift c | shift < 0 = uNode tchunk t shift c = chunk (get (uNode t) (digit c shift)) (shift - bits) c
index :: Array a -> Int -> aindex (Array t tl s) x | inBounds x s = if x >= treeSize s then fromU (get tl (x .&. mask)) else fromU (get (chunk t (rootShift (unsafeShiftR s bits)) (unsafeShiftR x bits)) (x .&. mask))index _ _ = error "TreeArray.index: out of bounds"
update :: Array a -> Int -> a -> Array aupdate (Array t tl s) x new | inBounds x s = if x >= treeSize s then Array t (set tl (x .&. mask) (toU new)) s else Array (updateGo (unsafeShiftR x bits) (x .&. mask) (toU new) t (rootShift (unsafeShiftR s bits))) tl supdate _ _ _ = error "TreeArray.update: out of bounds"
-- | Set element j of chunk c.updateGo :: Int -> Int -> U -> U -> Int -> UupdateGo c j new t shift | shift < 0 = nodeU (set (uNode t) j new)updateGo c j new t shift = let n = uNode t; i = digit c shift in nodeU (set n i (updateGo c j new (get n i) (shift - bits)))
{-# inline modify #-}modify :: forall a. Array a -> Int -> (a -> a) -> Array amodify (Array t tl s) x f | inBounds x s = if x >= treeSize s then Array t (leaf tl) s else Array (go t (rootShift (unsafeShiftR s bits))) tl s where j = x .&. mask c = unsafeShiftR x bits -- "f" stays captured, so that it is inlined here when modify is leaf :: Node -> Node leaf n = let x = fromU (get n j) in set n j (toU (f x)) go :: U -> Int -> U go t shift | shift < 0 = nodeU (leaf (uNode t)) go t shift = let n = uNode t; i = digit c shift in nodeU (set n i (go (get n i) (shift - bits)))modify _ _ _ = error "TreeArray.modify: out of bounds"
snoc :: Array a -> a -> Array asnoc (Array t tl s) new = let j = s .&. mask tl' = set tl j (toU new) in if j /= mask then Array t tl' (s + 1) -- the tail is full: push it into the tree. tl' is also reused as the new -- (logically empty) tail, since its slots are overwritten before being read. else Array (pushChunk t (unsafeShiftR s bits) (nodeU tl')) tl' (s + 1)
-- | Append a chunk to a tree of m chunks. The tree is full exactly when m is a-- power of W (including 1).pushChunk :: U -> Int -> U -> UpushChunk t m new | m .&. (m - 1) == 0 && m .&. powMask /= 0 = nodeU (node2 t (snocPath new (rootShift m))) | otherwise = snocGo m (countTrailingZeros m - bits) new t (rootShift (m + 1))
-- | A fresh path of nodes leading to "new" at index 0.snocPath :: U -> Int -> UsnocPath new shift | shift < 0 = newsnocPath new shift = nodeU (node1 (snocPath new (shift - bits)))
-- | Insert "new" at index m, given the shift at and below which subtrees are-- fresh (the low (shift + bits) bits of m are zero).snocGo :: Int -> Int -> U -> U -> Int -> UsnocGo m fresh new t shift | shift <= fresh = snocPath new shiftsnocGo m fresh new t shift = let n = uNode t; i = digit m shift in nodeU (set n i (snocGo m fresh new (get n i) (shift - bits)))
{-# inline foldr' #-}foldr' :: forall a b. (a -> b -> b) -> b -> Array a -> bfoldr' = \f b (Array t tl s) ->
let -- elements i, i-1 .. 0 of a node elems :: Node -> Int -> b -> b elems n i b | i < 0 = b elems n i b = let x = fromU (get n i) in elems n (i - 1) (f x b)
-- children i, i-1 .. 0 of a node fulls :: Node -> Int -> Int -> b -> b fulls n shift i b | i < 0 = b fulls n shift i b = fulls n shift (i - 1) (full (get n i) shift b)
notfull :: U -> Int -> Int -> b -> b notfull t shift m b | shift < 0 = elems (uNode t) mask b notfull t shift m b = let n = uNode t; shift' = shift - bits; ix = digit m shift in fulls n shift' (ix - 1) (notfull (get n ix) shift' m b)
full :: U -> Int -> b -> b full t shift b | shift < 0 = elems (uNode t) mask b full t shift b = fulls (uNode t) (shift - bits) mask b
m = unsafeShiftR s bits b' = elems tl ((s .&. mask) - 1) b
in case m of 0 -> b' _ -> notfull t (rootShift m) (m - 1) b'
toList :: Array a -> [a]toList = foldr' (:) []
instance Show a => Show (Array a) where show = show . toList