{- https://twitter.com/BrandonBloom/status/701261641971683328 https://github.com/clojure/clojure/blob/d5708425995e8c83157ad49007ec2f8f43d8eac8/src/jvm/clojure/lang/PersistentVector.java#L148-L164 -} {-@ LIQUID "--no-termination" @-} module PVec (arrayFor) where import Data.Bits -- | Simplified binary tree data Tree a = Leaf a | Node (Tree a) (Tree a) -- | Specify "height" of a tree {-@ measure height @-} height :: Tree a -> Int height (Leaf _) = 0 height (Node l _) = 1 + height l -- | A tree whose height is H {-@ type TreeH a H = {v:Tree | height v == H } @-} -- | Specify tree must be "balanced" {-@ data Tree a = Leaf a | Node { tLeft :: Tree a , tRight :: TreeH a (height tLeft) } @-} -- | Specify tree height is non-negative {-@ using (Tree a) as {v:Tree a | 0 <= height v} @-} -- | Vector stores the height data Vec a = Vec { vShift :: Int -- ^ height , vTree :: Tree a -- ^ actual nodes } -- | Refined type relates height of the `vTree` with `vShift` {-@ data Vec a = Vec { vShift :: Nat , vTree :: TreeH a vShift } @-} -------------------------------------------------------------------------------- arrayFor :: Int -> Vec a -> Maybe a arrayFor i (Vec l n) = loop l n where loop :: Int -> Tree a -> Maybe a loop level node | level > 0 = let b = shift i (- level) `mod` 2 -- get bit 0 or 1 node' = getNode node b -- get child level' = level - 1 -- next level in loop level' node' | otherwise = Just (getValue node) {-@ getNode :: t:{Tree a | 0 < height t} -> {v:Nat | v < 2} -> {v:Tree a | height v = height t - 1 } @-} getNode :: Tree a -> Int -> Tree a getNode (Node l _) 0 = l getNode (Node _ r) 1 = r getNode _ _ = impossible "provably safe" {-@ getValue :: {t:Tree a | 0 = height t} -> a @-} getValue :: Tree a -> a getValue (Leaf x) = x getValue _ = impossible "provably safe" {-@ impossible :: {v:String | false} -> a @-} impossible = error