Haskell / 16分木を用いた配列の実装

array1.hs

import Data.Bits

data Node =
    Branch { b0 :: Node, b1 :: Node, b2 :: Node, b3 :: Node,
        　    b4 :: Node, b5 :: Node, b6 :: Node, b7 :: Node,
　　         b8 :: Node, b9 :: Node, ba :: Node, bb :: Node,
　　         bc :: Node, bd :: Node, be :: Node, bf :: Node }
    | Leaf Int
    | Nil

showNode node = case node of
    Branch { 
        b0 = b0, b1 = b1, b2 = b2, b3 = b3,
        b4 = b4, b5 = b5, b6 = b6, b7 = b7,
        b8 = b8, b9 = b9, ba = ba, bb = bb,
        bc = bc, bd = bd, be = be, bf = bf } ->
            "Branch { 0: " ++ (showNode b0) ++ ", 1: " ++ (showNode b1)
            ++ ", 2: " ++ (showNode b2) ++ ", 3: " ++ (showNode b3)
            ++ ", 4: " ++ (showNode b4) ++ ", 5: " ++ (showNode b5)
            ++ ", 6: " ++ (showNode b6) ++ ", 7: " ++ (showNode b7)
            ++ ", 8: " ++ (showNode b8) ++ ", 9: " ++ (showNode b9)
            ++ ", 10: " ++ (showNode ba) ++ ", 11: " ++ (showNode bb)
            ++ ", 12: " ++ (showNode bc) ++ ", 13: " ++ (showNode bd)
            ++ ", 14: " ++ (showNode be) ++ ", 15: " ++ (showNode bf) ++ " }"
    Leaf n -> (show n)
    Nil -> "Nil"

newBranch = Branch { b0 = Nil, b1 = Nil, b2 = Nil, b3 = Nil,
                     b4 = Nil, b5 = Nil, b6 = Nil, b7 = Nil,
                     b8 = Nil, b9 = Nil, ba = Nil, bb = Nil,
                     bc = Nil, bd = Nil, be = Nil, bf = Nil }

data Array = Array { root :: Node }
--depth = 2 -- treats 0 ~ (2 ^ 8 - 1) = 0 ~ 255
depth = 5 -- treats 0 ~ (2 ^ 20 - 1) = 0 ~ 1048575

writeArray :: Array -> Int -> Int -> Array
writeArray ary idx val = Array $ writeArray' depth (root ary) idx val
writeArray' 0 node _ val = Leaf val
writeArray' i Nil idx val = case (idx .&. 15) of
    0  -> newBranch { b0 = writeArray' (i-1) Nil (shiftR idx 4) val }
    1  -> newBranch { b1 = writeArray' (i-1) Nil (shiftR idx 4) val }
    2  -> newBranch { b2 = writeArray' (i-1) Nil (shiftR idx 4) val }
    3  -> newBranch { b3 = writeArray' (i-1) Nil (shiftR idx 4) val }
    4  -> newBranch { b4 = writeArray' (i-1) Nil (shiftR idx 4) val }
    5  -> newBranch { b5 = writeArray' (i-1) Nil (shiftR idx 4) val }
    6  -> newBranch { b6 = writeArray' (i-1) Nil (shiftR idx 4) val }
    7  -> newBranch { b7 = writeArray' (i-1) Nil (shiftR idx 4) val }
    8  -> newBranch { b8 = writeArray' (i-1) Nil (shiftR idx 4) val }
    9  -> newBranch { b9 = writeArray' (i-1) Nil (shiftR idx 4) val }
    10 -> newBranch { ba = writeArray' (i-1) Nil (shiftR idx 4) val }
    11 -> newBranch { bb = writeArray' (i-1) Nil (shiftR idx 4) val }
    12 -> newBranch { bc = writeArray' (i-1) Nil (shiftR idx 4) val }
    13 -> newBranch { bd = writeArray' (i-1) Nil (shiftR idx 4) val }
    14 -> newBranch { be = writeArray' (i-1) Nil (shiftR idx 4) val }
    15 -> newBranch { bf = writeArray' (i-1) Nil (shiftR idx 4) val }
writeArray' i branch idx val = case (idx .&. 15) of
    0  -> branch { b0 = writeArray' (i-1) (b0 branch) (shiftR idx 4) val }
    1  -> branch { b1 = writeArray' (i-1) (b1 branch) (shiftR idx 4) val }
    2  -> branch { b2 = writeArray' (i-1) (b2 branch) (shiftR idx 4) val }
    3  -> branch { b3 = writeArray' (i-1) (b3 branch) (shiftR idx 4) val }
    4  -> branch { b4 = writeArray' (i-1) (b4 branch) (shiftR idx 4) val }
    5  -> branch { b5 = writeArray' (i-1) (b5 branch) (shiftR idx 4) val }
    6  -> branch { b6 = writeArray' (i-1) (b6 branch) (shiftR idx 4) val }
    7  -> branch { b7 = writeArray' (i-1) (b7 branch) (shiftR idx 4) val }
    8  -> branch { b8 = writeArray' (i-1) (b8 branch) (shiftR idx 4) val }
    9  -> branch { b9 = writeArray' (i-1) (b9 branch) (shiftR idx 4) val }
    10 -> branch { ba = writeArray' (i-1) (ba branch) (shiftR idx 4) val }
    11 -> branch { bb = writeArray' (i-1) (bb branch) (shiftR idx 4) val }
    12 -> branch { bc = writeArray' (i-1) (bc branch) (shiftR idx 4) val }
    13 -> branch { bd = writeArray' (i-1) (bd branch) (shiftR idx 4) val }
    14 -> branch { be = writeArray' (i-1) (be branch) (shiftR idx 4) val }
    15 -> branch { bf = writeArray' (i-1) (bf branch) (shiftR idx 4) val }

readArray :: Array -> Int -> Int
readArray ary idx = readArray' depth (root ary) idx
readArray' 0 (Leaf n) _ = n
readArray' 0 _ _ = error "invalid read"
readArray' i branch idx = case (idx .&. 15) of
    0  -> readArray' (i-1) (b0 branch) (shiftR idx 4)
    1  -> readArray' (i-1) (b1 branch) (shiftR idx 4)
    2  -> readArray' (i-1) (b2 branch) (shiftR idx 4)
    3  -> readArray' (i-1) (b3 branch) (shiftR idx 4)
    4  -> readArray' (i-1) (b4 branch) (shiftR idx 4)
    5  -> readArray' (i-1) (b5 branch) (shiftR idx 4)
    6  -> readArray' (i-1) (b6 branch) (shiftR idx 4)
    7  -> readArray' (i-1) (b7 branch) (shiftR idx 4)
    8  -> readArray' (i-1) (b8 branch) (shiftR idx 4)
    9  -> readArray' (i-1) (b9 branch) (shiftR idx 4)
    10 -> readArray' (i-1) (ba branch) (shiftR idx 4)
    11 -> readArray' (i-1) (bb branch) (shiftR idx 4)
    12 -> readArray' (i-1) (bc branch) (shiftR idx 4)
    13 -> readArray' (i-1) (bd branch) (shiftR idx 4)
    14 -> readArray' (i-1) (be branch) (shiftR idx 4)
    15 -> readArray' (i-1) (bf branch) (shiftR idx 4)

test1 = foldl simpleWriteRead (Array { root = newBranch }) [0..(2 ^ 20 - 1)]
    where
    simpleWriteRead ary n = (rd n) . (wr n) $ ary
    wr n ary = writeArray ary n n
    rd n ary = if (readArray ary n) /= n then error "fail test1" else ary

main = do
    return test1
    --putStrLn $ showNode $ root $ writeArray (Array { root = newBranch }) 2 5