    -----------------------------------------------------------------------------
    -- |
    -- Module    : Data.SBV.Tools.STree
    -- Copyright : (c) Levent Erkok
    -- License   : BSD3
    -- Maintainer: erkokl@gmail.com
    -- Stability : experimental
    --
    -- Implementation of full-binary symbolic trees, providing logarithmic
    -- time access to elements. Both reads and writes are supported.
    -----------------------------------------------------------------------------

    {-# LANGUAGE FlexibleContexts    #-}
    {-# LANGUAGE FlexibleInstances   #-}
    {-# LANGUAGE ScopedTypeVariables #-}
    {-# LANGUAGE TypeApplications    #-}

    {-# OPTIONS_GHC -Wall -Werror #-}

    module Data.SBV.Tools.STree (STree, readSTree, writeSTree, mkSTree) where

    import Data.SBV.Core.Data
    import Data.SBV.Core.Model

    import Data.Proxy

    -- | A symbolic tree containing values of type e, indexed by
    -- elements of type i. Note that these are full-trees, and their
    -- their shapes remain constant. There is no API provided that
    -- can change the shape of the tree. These structures are useful
    -- when dealing with data-structures that are indexed with symbolic
    -- values where access time is important. 'STree' structures provide
    -- logarithmic time reads and writes.
    type STree i e = STreeInternal (SBV i) (SBV e)

    -- Internal representation, not exposed to the user
    data STreeInternal i e = SLeaf e                        -- NB. parameter 'i' is phantom
                           | SBin  (STreeInternal i e) (STreeInternal i e)
                           deriving Int -> STreeInternal i e -> ShowS
    [STreeInternal i e] -> ShowS
    STreeInternal i e -> String
    (Int -> STreeInternal i e -> ShowS)
    -> (STreeInternal i e -> String)
    -> ([STreeInternal i e] -> ShowS)
    -> Show (STreeInternal i e)
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    forall i e. Show e => Int -> STreeInternal i e -> ShowS
    forall i e. Show e => [STreeInternal i e] -> ShowS
    forall i e. Show e => STreeInternal i e -> String
    $cshowsPrec :: forall i e. Show e => Int -> STreeInternal i e -> ShowS
    showsPrec :: Int -> STreeInternal i e -> ShowS
    $cshow :: forall i e. Show e => STreeInternal i e -> String
    show :: STreeInternal i e -> String
    $cshowList :: forall i e. Show e => [STreeInternal i e] -> ShowS
    showList :: [STreeInternal i e] -> ShowS
    Show

    instance SymVal e => Mergeable (STree i e) where
      symbolicMerge :: Bool -> SBool -> STree i e -> STree i e -> STree i e
    symbolicMerge Bool
    f SBool
    b (SLeaf SBV e
    i)  (SLeaf SBV e
    j)    = SBV e -> STree i e
    forall i e. e -> STreeInternal i e
    SLeaf (Bool -> SBool -> SBV e -> SBV e -> SBV e
    forall a. Mergeable a => Bool -> SBool -> a -> a -> a
    symbolicMerge Bool
    f SBool
    b SBV e
    i SBV e
    j)
      symbolicMerge Bool
    f SBool
    b (SBin STree i e
    l STree i e
    r) (SBin STree i e
    l' STree i e
    r') = STree i e -> STree i e -> STree i e
    forall i e.
    STreeInternal i e -> STreeInternal i e -> STreeInternal i e
    SBin  (Bool -> SBool -> STree i e -> STree i e -> STree i e
    forall a. Mergeable a => Bool -> SBool -> a -> a -> a
    symbolicMerge Bool
    f SBool
    b STree i e
    l STree i e
    l') (Bool -> SBool -> STree i e -> STree i e -> STree i e
    forall a. Mergeable a => Bool -> SBool -> a -> a -> a
    symbolicMerge Bool
    f SBool
    b STree i e
    r STree i e
    r')
      symbolicMerge Bool
    _ SBool
    _ STree i e
    _          STree i e
    _            = String -> STree i e
    forall a. HasCallStack => String -> a
    error String
    "SBV.STree.symbolicMerge: Impossible happened while merging states"

    -- | Reading a value. We bit-blast the index and descend down the full tree
    -- according to bit-values.
    readSTree :: (SFiniteBits i, SymVal e) => STree i e -> SBV i -> SBV e
    readSTree :: forall i e.
    (SFiniteBits i, SymVal e) =>
    STree i e -> SBV i -> SBV e
    readSTree STree i e
    s SBV i
    i = [SBool] -> STree i e -> SBV e
    forall {e} {i}. Mergeable e => [SBool] -> STreeInternal i e -> e
    walk (SBV i -> [SBool]
    forall a. SFiniteBits a => SBV a -> [SBool]
    blastBE SBV i
    i) STree i e
    s
      where walk :: [SBool] -> STreeInternal i e -> e
    walk []     (SLeaf e
    v)  = e
    v
            walk (SBool
    b:[SBool]
    bs) (SBin STreeInternal i e
    l STreeInternal i e
    r) = SBool -> e -> e -> e
    forall a. Mergeable a => SBool -> a -> a -> a
    ite SBool
    b ([SBool] -> STreeInternal i e -> e
    walk [SBool]
    bs STreeInternal i e
    r) ([SBool] -> STreeInternal i e -> e
    walk [SBool]
    bs STreeInternal i e
    l)
            walk [SBool]
    _      STreeInternal i e
    _          = String -> e
    forall a. HasCallStack => String -> a
    error (String -> e) -> String -> e
    forall a b. (a -> b) -> a -> b
    $ String
    "SBV.STree.readSTree: Impossible happened while reading: " String -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ SBV i -> String
    forall a. Show a => a -> String
    show SBV i
    i

    -- | Writing a value, similar to how reads are done. The important thing is that the tree
    -- representation keeps updates to a minimum.
    writeSTree :: (SFiniteBits i, SymVal e) => STree i e -> SBV i -> SBV e -> STree i e
    writeSTree :: forall i e.
    (SFiniteBits i, SymVal e) =>
    STree i e -> SBV i -> SBV e -> STree i e
    writeSTree STreeInternal (SBV i) (SBV e)
    s SBV i
    i SBV e
    j = [SBool]
    -> STreeInternal (SBV i) (SBV e) -> STreeInternal (SBV i) (SBV e)
    forall {i}.
    Mergeable (STreeInternal i (SBV e)) =>
    [SBool] -> STreeInternal i (SBV e) -> STreeInternal i (SBV e)
    walk (SBV i -> [SBool]
    forall a. SFiniteBits a => SBV a -> [SBool]
    blastBE SBV i
    i) STreeInternal (SBV i) (SBV e)
    s
      where walk :: [SBool] -> STreeInternal i (SBV e) -> STreeInternal i (SBV e)
    walk []     STreeInternal i (SBV e)
    _          = SBV e -> STreeInternal i (SBV e)
    forall i e. e -> STreeInternal i e
    SLeaf SBV e
    j
            walk (SBool
    b:[SBool]
    bs) (SBin STreeInternal i (SBV e)
    l STreeInternal i (SBV e)
    r) = STreeInternal i (SBV e)
    -> STreeInternal i (SBV e) -> STreeInternal i (SBV e)
    forall i e.
    STreeInternal i e -> STreeInternal i e -> STreeInternal i e
    SBin (SBool
    -> STreeInternal i (SBV e)
    -> STreeInternal i (SBV e)
    -> STreeInternal i (SBV e)
    forall a. Mergeable a => SBool -> a -> a -> a
    ite SBool
    b STreeInternal i (SBV e)
    l ([SBool] -> STreeInternal i (SBV e) -> STreeInternal i (SBV e)
    walk [SBool]
    bs STreeInternal i (SBV e)
    l)) (SBool
    -> STreeInternal i (SBV e)
    -> STreeInternal i (SBV e)
    -> STreeInternal i (SBV e)
    forall a. Mergeable a => SBool -> a -> a -> a
    ite SBool
    b ([SBool] -> STreeInternal i (SBV e) -> STreeInternal i (SBV e)
    walk [SBool]
    bs STreeInternal i (SBV e)
    r) STreeInternal i (SBV e)
    r)
            walk [SBool]
    _      STreeInternal i (SBV e)
    _          = String -> STreeInternal i (SBV e)
    forall a. HasCallStack => String -> a
    error (String -> STreeInternal i (SBV e))
    -> String -> STreeInternal i (SBV e)
    forall a b. (a -> b) -> a -> b
    $ String
    "SBV.STree.writeSTree: Impossible happened while reading: " String -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ SBV i -> String
    forall a. Show a => a -> String
    show SBV i
    i

    -- | Construct the fully balanced initial tree using the given values.
    mkSTree :: forall i e. HasKind i => [SBV e] -> STree i e
    mkSTree :: forall i e. HasKind i => [SBV e] -> STree i e
    mkSTree [SBV e]
    ivals
      | Proxy i -> Bool
    forall a. HasKind a => a -> Bool
    isReal (forall t. Proxy t
    forall {k} (t :: k). Proxy t
    Proxy @i)
      = String -> STree i e
    forall a. HasCallStack => String -> a
    error String
    "SBV.STree.mkSTree: Cannot build a real-valued sized tree"
      | Bool -> Bool
    not (Proxy i -> Bool
    forall a. HasKind a => a -> Bool
    isBounded (forall t. Proxy t
    forall {k} (t :: k). Proxy t
    Proxy @i))
      = String -> STree i e
    forall a. HasCallStack => String -> a
    error String
    "SBV.STree.mkSTree: Cannot build an infinitely large tree"
      | Int
    reqd Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    /= Int
    given
      = String -> STree i e
    forall a. HasCallStack => String -> a
    error (String -> STree i e) -> String -> STree i e
    forall a b. (a -> b) -> a -> b
    $ String
    "SBV.STree.mkSTree: Required " String -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ Int -> String
    forall a. Show a => a -> String
    show Int
    reqd String -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ String
    " elements, received: " String -> ShowS
    forall a. [a] -> [a] -> [a]
    ++ Int -> String
    forall a. Show a => a -> String
    show Int
    given
      | Bool
    True
      = [SBV e] -> STree i e
    forall {e} {i}. [e] -> STreeInternal i e
    go [SBV e]
    ivals
      where reqd :: Int
    reqd = Int
    2 Int -> Int -> Int
    forall a b. (Num a, Integral b) => a -> b -> a
    ^ Proxy i -> Int
    forall a. HasKind a => a -> Int
    intSizeOf (forall t. Proxy t
    forall {k} (t :: k). Proxy t
    Proxy @i)
            given :: Int
    given = [SBV e] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SBV e]
    ivals
            go :: [e] -> STreeInternal i e
    go []  = String -> STreeInternal i e
    forall a. HasCallStack => String -> a
    error String
    "SBV.STree.mkSTree: Impossible happened, ran out of elements"
            go [e
    l] = e -> STreeInternal i e
    forall i e. e -> STreeInternal i e
    SLeaf e
    l
            go [e]
    ns  = let ([e]
    l, [e]
    r) = Int -> [e] -> ([e], [e])
    forall a. Int -> [a] -> ([a], [a])
    splitAt ([e] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [e]
    ns Int -> Int -> Int
    forall a. Integral a => a -> a -> a
    `div` Int
    2) [e]
    ns in STreeInternal i e -> STreeInternal i e -> STreeInternal i e
    forall i e.
    STreeInternal i e -> STreeInternal i e -> STreeInternal i e
    SBin ([e] -> STreeInternal i e
    go [e]
    l) ([e] -> STreeInternal i e
    go [e]
    r)