-----------------------------------------------------------------------------
-- |
-- 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
showList :: [STreeInternal i e] -> ShowS
$cshowList :: forall i e. Show e => [STreeInternal i e] -> ShowS
show :: STreeInternal i e -> String
$cshow :: forall i e. Show e => STreeInternal i e -> String
showsPrec :: Int -> STreeInternal i e -> ShowS
$cshowsPrec :: forall i e. Show e => Int -> 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 :: STree i e -> SBV i -> SBV e
readSTree STree i e
s SBV i
i = [SBool] -> STree i e -> SBV e
forall p i. Mergeable p => [SBool] -> STreeInternal i p -> p
walk (SBV i -> [SBool]
forall a. SFiniteBits a => SBV a -> [SBool]
blastBE SBV i
i) STree i e
s
  where walk :: [SBool] -> STreeInternal i p -> p
walk []     (SLeaf p
v)  = p
v
        walk (SBool
b:[SBool]
bs) (SBin STreeInternal i p
l STreeInternal i p
r) = SBool -> p -> p -> p
forall a. Mergeable a => SBool -> a -> a -> a
ite SBool
b ([SBool] -> STreeInternal i p -> p
walk [SBool]
bs STreeInternal i p
r) ([SBool] -> STreeInternal i p -> p
walk [SBool]
bs STreeInternal i p
l)
        walk [SBool]
_      STreeInternal i p
_          = String -> p
forall a. HasCallStack => String -> a
error (String -> p) -> String -> p
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 :: STree i e -> SBV i -> SBV e -> STree i e
writeSTree STree i e
s SBV i
i SBV e
j = [SBool] -> STree i e -> STree i 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) STree i 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 :: [SBV e] -> STree i e
mkSTree [SBV e]
ivals
  | Proxy i -> Bool
forall a. HasKind a => a -> Bool
isReal (Proxy i
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 (Proxy i
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 (Proxy i
forall k (t :: k). Proxy t
Proxy @i)
        given :: Int
given = [SBV e] -> 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 (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)