Commits

Bryan O'Sullivan committed 32c7576

A little more flesh on the bones

Comments (0)

Files changed (3)

     , fromList
     , fromVector
     , index
+    , isSane
     , snoc
     , update
+    , Zipper
+    , toZipper
+    , snocZ
     ) where
 
-import Debug.Trace
 import Control.DeepSeq
 import Data.Bits hiding (shift)
+import Data.Vector (Vector)
 import qualified Data.Vector as V
-import Data.Vector (Vector)
+import Debug.Trace
 
 data Vect a = Node !Int !(Vector (Vect a))
             | Leaf !(Vector a)
               deriving (Eq, Ord)
 
+instance NFData a => NFData (V.Vector a) where
+    rnf = V.foldl' (const rnf) ()
+
 instance NFData a => NFData (Vect a) where
-    rnf (Leaf v) = V.foldl' (const rnf) () v
-    rnf (Node _ v) = V.foldl' (const rnf) () v
+    rnf (Leaf v) = rnf v
+    rnf (Node _ v) = rnf v
 
 instance Show a => Show (Vect a) where
     show (Leaf v)   = show (V.toList v)
                              (map (Node h . V.fromList) $ chunksOf factor ns)
     where h' = h+shift
 
+empty :: Vect a
+empty = Leaf V.empty
+
 chunksOf :: Int -> [a] -> [[a]]
 chunksOf k = go
   where go [] = []
 shiftOf (Leaf _)   = 0
 shiftOf (Node s _) = s
 
+isFull :: Vect a -> Bool
+isFull (Leaf v)   = V.length v == factor
+isFull (Node s v) = V.length v == factor && V.all isFull v &&
+                    V.all ((==s-shift) . shiftOf) v
+
+isSane :: Vect a -> Bool
+isSane (Leaf v)   = V.length v <= factor
+isSane (Node s v) = V.length v <= factor &&
+                    V.all isFull (V.init v) &&
+                    isSane (V.last v) &&
+                    ((==s-shift) . shiftOf) (V.last v)
+
+
 snoc :: Vect a -> a -> Vect a
 snoc t n = case go t of
              Left n'  -> n'
             Left n'                 -> Left $! Node s (V.init v `V.snoc` n')
             Right n'
               | V.length v < factor -> Left $! Node s (v `V.snoc` n')
-              | otherwise           -> Right $! Node (shift+s) (V.singleton n')
+              | otherwise           -> Right $! Node s (V.singleton n')
+
+snocChunk :: Vect a -> Vector a -> Vect a
+snocChunk t c = case go t of
+                  Left n' -> n'
+                  Right n' -> Node (shift+shiftOf n') (V.fromList [t,n'])
+  where
+    go (Leaf _)    = Right l
+    go (Node s v)
+      | s == shift = if V.length v < factor
+                     then Left $! Node s (v `V.snoc` l)
+                     else Right $! Node (shift+s) (V.singleton l)
+      | otherwise  = case go (V.last v) of
+                       Left n' -> Left $! Node s (V.init v `V.snoc` n')
+                       Right n'
+                         | V.length v < factor -> Left $! Node s (v `V.snoc` n')
+                         | otherwise           -> Right $! Node s (V.singleton n')
+    l = Leaf c
+
+data Zipper a = Z !(Vect a) !Int !(Vector a)
+
+instance Show a => Show (Zipper a) where
+    show (Z v k c) = "Z " ++ show v ++ " " ++ show k ++ " " ++ show (V.toList c)
+
+instance NFData a => NFData (Zipper a) where
+    rnf (Z v _ c) = rnf v `seq` rnf c `seq` ()
+
+toZipper :: Vect a -> Int -> Zipper a
+toZipper t k = go t
+  where go (Leaf v)   = Z t (k - (k .&. mask)) v
+        go (Node s v) = go (v V.! ((k `shiftR` s) .&. mask))
+
+fromZipper :: Zipper a -> Vect a
+fromZipper (Z t k c) = go t
+  where go (Leaf _)   = Leaf c
+        go (Node s v) = Node s (v V.// [(j, go (v V.! j))])
+            where j = (k `shiftR` s) .&. mask
+
+updateZ :: Zipper a -> Int -> a -> Zipper a
+updateZ z@(Z t k c) j n
+    | j-jm == k = Z t k (c V.// [(jm,n)])
+    | otherwise = updateZ (toZipper (fromZipper z) j) j n
+    where jm = j .&. mask
+
+snocZ :: Zipper a -> a -> Zipper a
+snocZ (Z v k c) n | V.length c < factor = Z v k (V.snoc c n)
+                  | otherwise = Z (snocChunk v c) k (V.singleton n)
 
 mapVect :: (a -> b) -> Vect a -> Vect b
 mapVect f = go

tests/Properties.hs

+import Data.Vectro
+import Test.Framework (defaultMain, testGroup)
+import Test.QuickCheck
+import Test.Framework.Providers.QuickCheck2 (testProperty)
+
+ints :: Vect Int -> Vect Int
+ints a = a
+
+instance Arbitrary a => Arbitrary (Vect a) where
+    arbitrary = fromList `fmap` arbitrary
+
+t_isSane = isSane . ints
+
+t_snoc_sane a = isSane . (`snoc` a) . ints
+
+main = defaultMain tests
+
+tests = [
+   testProperty "t_isSane" t_isSane,
+   testProperty "t_snoc_sane" t_snoc_sane
+ ]

tests/benchmarks/Benchmark.hs

 import Data.Vector as V
 
 main = do
-  let l = [1..(1000::Int)]
+  let l = [1..(10000::Int)]
   let v  = V.fromList l
   let tv = TV.fromVector v
+  let vec = TV.fromVect tv
   defaultMain [
         bgroup "build" [
             bench "tv" $ nf TV.fromList l
           , bench "v" $ whnf V.fromList l
           , bench "tvv" $ nf TV.fromVector v
           ],
+        bgroup "snoc" [
+            bench "v" $ nf (V.foldl' V.snoc v) v
+          , bench "tv" $ nf (V.foldl' TV.snoc tv) v
+          , bench "vec" $ nf (V.foldl' TV.snocV vec) v
+          ],
         bgroup "update" [
             bench "v" $ whnf (V.foldl' (\v i -> v V.// [(i-1,0)]) v) v
           , bench "tv" $ nf (V.foldl' (\t i -> TV.update t (i-1) 0) tv) v