{-# LANGUAGE CPP #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE Trustworthy #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise -fplugin GHC.TypeLits.KnownNat.Solver #-}
module Clash.Sized.RTree
(
RTree (LR, BR)
, treplicate
, trepeat
, indexTree
, tindices
, replaceTree
, tmap
, tzipWith
, tzip
, tunzip
, tfold
, tdfold
, v2t
, t2v
, lazyT
)
where
import Control.Applicative (liftA2)
import Control.DeepSeq (NFData(..))
import qualified Control.Lens as Lens
import Data.Default.Class (Default (..))
import Data.Either (isLeft)
import Data.Foldable (toList)
import Data.Kind (Type)
import Data.Singletons.Prelude (Apply, TyFun, type (@@))
import Data.Proxy (Proxy (..))
import GHC.TypeLits (KnownNat, Nat, type (+), type (^), type (*))
import Language.Haskell.TH.Syntax (Lift(..))
import Prelude hiding ((++), (!!))
import Test.QuickCheck (Arbitrary (..), CoArbitrary (..))
import Clash.Class.BitPack (BitPack (..), packXWith)
import Clash.Promoted.Nat (SNat (..), UNat (..), pow2SNat, snatToNum,
subSNat, toUNat)
import Clash.Promoted.Nat.Literals (d1)
import Clash.Sized.Index (Index)
import Clash.Sized.Vector (Vec (..), (!!), (++), dtfold, replace)
import Clash.XException
(ShowX (..), NFDataX (..), isX, showsX, showsPrecXWith)
data RTree :: Nat -> Type -> Type where
LR_ :: a -> RTree 0 a
BR_ :: RTree d a -> RTree d a -> RTree (d+1) a
instance NFData a => NFData (RTree d a) where
rnf :: RTree d a -> ()
rnf (LR_ x :: a
x) = a -> ()
forall a. NFData a => a -> ()
rnf a
x
rnf (BR_ l :: RTree d a
l r :: RTree d a
r ) = RTree d a -> ()
forall a. NFData a => a -> ()
rnf RTree d a
l () -> () -> ()
forall a b. a -> b -> b
`seq` RTree d a -> ()
forall a. NFData a => a -> ()
rnf RTree d a
r
textract :: RTree 0 a -> a
(LR_ x :: a
x) = a
x
textract (BR_ _ _) = [Char] -> a
forall a. HasCallStack => [Char] -> a
error ([Char] -> a) -> [Char] -> a
forall a b. (a -> b) -> a -> b
$ "textract: nodes hold no values"
{-# NOINLINE textract #-}
tsplit :: RTree (d+1) a -> (RTree d a,RTree d a)
tsplit :: RTree (d + 1) a -> (RTree d a, RTree d a)
tsplit (BR_ l :: RTree d a
l r :: RTree d a
r) = (RTree d a
RTree d a
l,RTree d a
RTree d a
r)
tsplit (LR_ _) = [Char] -> (RTree d a, RTree d a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (RTree d a, RTree d a))
-> [Char] -> (RTree d a, RTree d a)
forall a b. (a -> b) -> a -> b
$ "tsplit: leaf is atomic"
{-# NOINLINE tsplit #-}
pattern LR :: a -> RTree 0 a
pattern $bLR :: a -> RTree 0 a
$mLR :: forall r a. RTree 0 a -> (a -> r) -> (Void# -> r) -> r
LR x <- (textract -> x)
where
LR x :: a
x = a -> RTree 0 a
forall a. a -> RTree 0 a
LR_ a
x
pattern BR :: RTree d a -> RTree d a -> RTree (d+1) a
pattern $bBR :: RTree d a -> RTree d a -> RTree (d + 1) a
$mBR :: forall r (d :: Nat) a.
RTree (d + 1) a
-> (RTree d a -> RTree d a -> r) -> (Void# -> r) -> r
BR l r <- ((\t -> (tsplit t)) -> (l,r))
where
BR l :: RTree d a
l r :: RTree d a
r = RTree d a -> RTree d a -> RTree (d + 1) a
forall (d :: Nat) a. RTree d a -> RTree d a -> RTree (d + 1) a
BR_ RTree d a
l RTree d a
r
instance (KnownNat d, Eq a) => Eq (RTree d a) where
== :: RTree d a -> RTree d a -> Bool
(==) t1 :: RTree d a
t1 t2 :: RTree d a
t2 = Vec (2 ^ d) a -> Vec (2 ^ d) a -> Bool
forall a. Eq a => a -> a -> Bool
(==) (RTree d a -> Vec (2 ^ d) a
forall (d :: Nat) a. KnownNat d => RTree d a -> Vec (2 ^ d) a
t2v RTree d a
t1) (RTree d a -> Vec (2 ^ d) a
forall (d :: Nat) a. KnownNat d => RTree d a -> Vec (2 ^ d) a
t2v RTree d a
t2)
instance (KnownNat d, Ord a) => Ord (RTree d a) where
compare :: RTree d a -> RTree d a -> Ordering
compare t1 :: RTree d a
t1 t2 :: RTree d a
t2 = Vec (2 ^ d) a -> Vec (2 ^ d) a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (RTree d a -> Vec (2 ^ d) a
forall (d :: Nat) a. KnownNat d => RTree d a -> Vec (2 ^ d) a
t2v RTree d a
t1) (RTree d a -> Vec (2 ^ d) a
forall (d :: Nat) a. KnownNat d => RTree d a -> Vec (2 ^ d) a
t2v RTree d a
t2)
instance Show a => Show (RTree n a) where
showsPrec :: Int -> RTree n a -> ShowS
showsPrec _ (LR_ a :: a
a) = a -> ShowS
forall a. Show a => a -> ShowS
shows a
a
showsPrec _ (BR_ l :: RTree d a
l r :: RTree d a
r) = \s :: [Char]
s -> '<'Char -> ShowS
forall a. a -> [a] -> [a]
:RTree d a -> ShowS
forall a. Show a => a -> ShowS
shows RTree d a
l (','Char -> ShowS
forall a. a -> [a] -> [a]
:RTree d a -> ShowS
forall a. Show a => a -> ShowS
shows RTree d a
r ('>'Char -> ShowS
forall a. a -> [a] -> [a]
:[Char]
s))
instance ShowX a => ShowX (RTree n a) where
showsPrecX :: Int -> RTree n a -> ShowS
showsPrecX = (Int -> RTree n a -> ShowS) -> Int -> RTree n a -> ShowS
forall a. (Int -> a -> ShowS) -> Int -> a -> ShowS
showsPrecXWith Int -> RTree n a -> ShowS
forall (d :: Nat). Int -> RTree d a -> ShowS
go
where
go :: Int -> RTree d a -> ShowS
go :: Int -> RTree d a -> ShowS
go _ (LR_ a :: a
a) = a -> ShowS
forall a. ShowX a => a -> ShowS
showsX a
a
go _ (BR_ l :: RTree d a
l r :: RTree d a
r) = \s :: [Char]
s -> '<'Char -> ShowS
forall a. a -> [a] -> [a]
:RTree d a -> ShowS
forall a. ShowX a => a -> ShowS
showsX RTree d a
l (','Char -> ShowS
forall a. a -> [a] -> [a]
:RTree d a -> ShowS
forall a. ShowX a => a -> ShowS
showsX RTree d a
r ('>'Char -> ShowS
forall a. a -> [a] -> [a]
:[Char]
s))
instance KnownNat d => Functor (RTree d) where
fmap :: (a -> b) -> RTree d a -> RTree d b
fmap = (a -> b) -> RTree d a -> RTree d b
forall (d :: Nat) a b.
KnownNat d =>
(a -> b) -> RTree d a -> RTree d b
tmap
instance KnownNat d => Applicative (RTree d) where
pure :: a -> RTree d a
pure = a -> RTree d a
forall (d :: Nat) a. KnownNat d => a -> RTree d a
trepeat
<*> :: RTree d (a -> b) -> RTree d a -> RTree d b
(<*>) = ((a -> b) -> a -> b) -> RTree d (a -> b) -> RTree d a -> RTree d b
forall a b c (d :: Nat).
KnownNat d =>
(a -> b -> c) -> RTree d a -> RTree d b -> RTree d c
tzipWith (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($)
instance KnownNat d => Foldable (RTree d) where
foldMap :: (a -> m) -> RTree d a -> m
foldMap f :: a -> m
f = (a -> m) -> (m -> m -> m) -> RTree d a -> m
forall (d :: Nat) a b.
KnownNat d =>
(a -> b) -> (b -> b -> b) -> RTree d a -> b
tfold a -> m
f m -> m -> m
forall a. Monoid a => a -> a -> a
mappend
data TraversableTree (g :: Type -> Type) (a :: Type) (f :: TyFun Nat Type) :: Type
type instance Apply (TraversableTree f a) d = f (RTree d a)
instance KnownNat d => Traversable (RTree d) where
traverse :: forall f a b . Applicative f => (a -> f b) -> RTree d a -> f (RTree d b)
traverse :: (a -> f b) -> RTree d a -> f (RTree d b)
traverse f :: a -> f b
f = Proxy (TraversableTree f b)
-> (a -> TraversableTree f b @@ 0)
-> (forall (l :: Nat).
SNat l
-> (TraversableTree f b @@ l)
-> (TraversableTree f b @@ l)
-> TraversableTree f b @@ (l + 1))
-> RTree d a
-> TraversableTree f b @@ d
forall (p :: TyFun Nat Type -> Type) (k :: Nat) a.
KnownNat k =>
Proxy p
-> (a -> p @@ 0)
-> (forall (l :: Nat).
SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1))
-> RTree k a
-> p @@ k
tdfold (Proxy (TraversableTree f b)
forall k (t :: k). Proxy t
Proxy @(TraversableTree f b))
((b -> RTree 0 b) -> f b -> f (RTree 0 b)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> RTree 0 b
forall a. a -> RTree 0 a
LR (f b -> f (RTree 0 b)) -> (a -> f b) -> a -> f (RTree 0 b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> f b
f)
((f (RTree l b) -> f (RTree l b) -> f (RTree (l + 1) b))
-> SNat l -> f (RTree l b) -> f (RTree l b) -> f (RTree (l + 1) b)
forall a b. a -> b -> a
const ((RTree l b -> RTree l b -> RTree (l + 1) b)
-> f (RTree l b) -> f (RTree l b) -> f (RTree (l + 1) b)
forall (f :: Type -> Type) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 RTree l b -> RTree l b -> RTree (l + 1) b
forall (d :: Nat) a. RTree d a -> RTree d a -> RTree (d + 1) a
BR))
instance (KnownNat d, BitPack a) =>
BitPack (RTree d a) where
type BitSize (RTree d a) = (2^d) * (BitSize a)
pack :: RTree d a -> BitVector (BitSize (RTree d a))
pack = (RTree d a -> BitVector ((2 ^ d) * BitSize a))
-> RTree d a -> BitVector ((2 ^ d) * BitSize a)
forall (n :: Nat) a.
KnownNat n =>
(a -> BitVector n) -> a -> BitVector n
packXWith (Vec (2 ^ d) a -> BitVector ((2 ^ d) * BitSize a)
forall a. BitPack a => a -> BitVector (BitSize a)
pack (Vec (2 ^ d) a -> BitVector ((2 ^ d) * BitSize a))
-> (RTree d a -> Vec (2 ^ d) a)
-> RTree d a
-> BitVector ((2 ^ d) * BitSize a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RTree d a -> Vec (2 ^ d) a
forall (d :: Nat) a. KnownNat d => RTree d a -> Vec (2 ^ d) a
t2v)
unpack :: BitVector (BitSize (RTree d a)) -> RTree d a
unpack = Vec (2 ^ d) a -> RTree d a
forall (d :: Nat) a. KnownNat d => Vec (2 ^ d) a -> RTree d a
v2t (Vec (2 ^ d) a -> RTree d a)
-> (BitVector ((2 ^ d) * BitSize a) -> Vec (2 ^ d) a)
-> BitVector ((2 ^ d) * BitSize a)
-> RTree d a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitVector ((2 ^ d) * BitSize a) -> Vec (2 ^ d) a
forall a. BitPack a => BitVector (BitSize a) -> a
unpack
type instance Lens.Index (RTree d a) = Int
type instance Lens.IxValue (RTree d a) = a
instance KnownNat d => Lens.Ixed (RTree d a) where
ix :: Index (RTree d a) -> Traversal' (RTree d a) (IxValue (RTree d a))
ix i :: Index (RTree d a)
i f :: IxValue (RTree d a) -> f (IxValue (RTree d a))
f t :: RTree d a
t = Int -> a -> RTree d a -> RTree d a
forall (d :: Nat) i a.
(KnownNat d, Enum i) =>
i -> a -> RTree d a -> RTree d a
replaceTree Int
Index (RTree d a)
i (a -> RTree d a -> RTree d a) -> f a -> f (RTree d a -> RTree d a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxValue (RTree d a) -> f (IxValue (RTree d a))
f (RTree d a -> Int -> a
forall (d :: Nat) i a. (KnownNat d, Enum i) => RTree d a -> i -> a
indexTree RTree d a
t Int
Index (RTree d a)
i) f (RTree d a -> RTree d a) -> f (RTree d a) -> f (RTree d a)
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> RTree d a -> f (RTree d a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure RTree d a
t
instance (KnownNat d, Default a) => Default (RTree d a) where
def :: RTree d a
def = a -> RTree d a
forall (d :: Nat) a. KnownNat d => a -> RTree d a
trepeat a
forall a. Default a => a
def
instance Lift a => Lift (RTree d a) where
lift :: RTree d a -> Q Exp
lift (LR_ a :: a
a) = [| LR_ a |]
lift (BR_ t1 :: RTree d a
t1 t2 :: RTree d a
t2) = [| BR_ $(lift t1) $(lift t2) |]
instance (KnownNat d, Arbitrary a) => Arbitrary (RTree d a) where
arbitrary :: Gen (RTree d a)
arbitrary = RTree d (Gen a) -> Gen (RTree d a)
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (Gen a -> RTree d (Gen a)
forall (d :: Nat) a. KnownNat d => a -> RTree d a
trepeat Gen a
forall a. Arbitrary a => Gen a
arbitrary)
shrink :: RTree d a -> [RTree d a]
shrink = RTree d [a] -> [RTree d a]
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (RTree d [a] -> [RTree d a])
-> (RTree d a -> RTree d [a]) -> RTree d a -> [RTree d a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> [a]) -> RTree d a -> RTree d [a]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> [a]
forall a. Arbitrary a => a -> [a]
shrink
instance (KnownNat d, CoArbitrary a) => CoArbitrary (RTree d a) where
coarbitrary :: RTree d a -> Gen b -> Gen b
coarbitrary = [a] -> Gen b -> Gen b
forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary ([a] -> Gen b -> Gen b)
-> (RTree d a -> [a]) -> RTree d a -> Gen b -> Gen b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RTree d a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList
instance (KnownNat d, NFDataX a) => NFDataX (RTree d a) where
deepErrorX :: [Char] -> RTree d a
deepErrorX x :: [Char]
x = a -> RTree d a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ([Char] -> a
forall a. (NFDataX a, HasCallStack) => [Char] -> a
deepErrorX [Char]
x)
rnfX :: RTree d a -> ()
rnfX t :: RTree d a
t = if Either [Char] (RTree d a) -> Bool
forall a b. Either a b -> Bool
isLeft (RTree d a -> Either [Char] (RTree d a)
forall a. a -> Either [Char] a
isX RTree d a
t) then () else RTree d a -> ()
go RTree d a
t
where
go :: RTree d a -> ()
go :: RTree d a -> ()
go (LR_ x :: a
x) = a -> ()
forall a. NFDataX a => a -> ()
rnfX a
x
go (BR_ l :: RTree d a
l r :: RTree d a
r) = RTree d a -> ()
forall a. NFDataX a => a -> ()
rnfX RTree d a
l () -> () -> ()
forall a b. a -> b -> b
`seq` RTree d a -> ()
forall a. NFDataX a => a -> ()
rnfX RTree d a
r
hasUndefined :: RTree d a -> Bool
hasUndefined t :: RTree d a
t = if Either [Char] (RTree d a) -> Bool
forall a b. Either a b -> Bool
isLeft (RTree d a -> Either [Char] (RTree d a)
forall a. a -> Either [Char] a
isX RTree d a
t) then Bool
True else RTree d a -> Bool
go RTree d a
t
where
go :: RTree d a -> Bool
go :: RTree d a -> Bool
go (LR_ x :: a
x) = a -> Bool
forall a. NFDataX a => a -> Bool
hasUndefined a
x
go (BR_ l :: RTree d a
l r :: RTree d a
r) = RTree d a -> Bool
forall a. NFDataX a => a -> Bool
hasUndefined RTree d a
l Bool -> Bool -> Bool
|| RTree d a -> Bool
forall a. NFDataX a => a -> Bool
hasUndefined RTree d a
r
ensureSpine :: RTree d a -> RTree d a
ensureSpine = (a -> a) -> RTree d a -> RTree d a
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. NFDataX a => a -> a
ensureSpine (RTree d a -> RTree d a)
-> (RTree d a -> RTree d a) -> RTree d a -> RTree d a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RTree d a -> RTree d a
forall (d :: Nat) a. KnownNat d => RTree d a -> RTree d a
lazyT
tdfold :: forall p k a . KnownNat k
=> Proxy (p :: TyFun Nat Type -> Type)
-> (a -> (p @@ 0))
-> (forall l . SNat l -> (p @@ l) -> (p @@ l) -> (p @@ (l+1)))
-> RTree k a
-> (p @@ k)
tdfold :: Proxy p
-> (a -> p @@ 0)
-> (forall (l :: Nat).
SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1))
-> RTree k a
-> p @@ k
tdfold _ f :: a -> p @@ 0
f g :: forall (l :: Nat). SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1)
g = SNat k -> RTree k a -> p @@ k
forall (m :: Nat). SNat m -> RTree m a -> p @@ m
go SNat k
forall (n :: Nat). KnownNat n => SNat n
SNat
where
go :: SNat m -> RTree m a -> (p @@ m)
go :: SNat m -> RTree m a -> p @@ m
go _ (LR_ a :: a
a) = a -> p @@ 0
f a
a
go sn :: SNat m
sn (BR_ l :: RTree d a
l r :: RTree d a
r) = let sn' :: SNat d
sn' = SNat m
SNat (d + 1)
sn SNat (d + 1) -> SNat 1 -> SNat d
forall (a :: Nat) (b :: Nat). SNat (a + b) -> SNat b -> SNat a
`subSNat` SNat 1
d1
in SNat d -> (p @@ d) -> (p @@ d) -> p @@ (d + 1)
forall (l :: Nat). SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1)
g SNat d
sn' (SNat d -> RTree d a -> p @@ d
forall (m :: Nat). SNat m -> RTree m a -> p @@ m
go SNat d
sn' RTree d a
l) (SNat d -> RTree d a -> p @@ d
forall (m :: Nat). SNat m -> RTree m a -> p @@ m
go SNat d
sn' RTree d a
r)
{-# NOINLINE tdfold #-}
data TfoldTree (a :: Type) (f :: TyFun Nat Type) :: Type
type instance Apply (TfoldTree a) d = a
tfold :: forall d a b .
KnownNat d
=> (a -> b)
-> (b -> b -> b)
-> RTree d a
-> b
tfold :: (a -> b) -> (b -> b -> b) -> RTree d a -> b
tfold f :: a -> b
f g :: b -> b -> b
g = Proxy (TfoldTree b)
-> (a -> TfoldTree b @@ 0)
-> (forall (l :: Nat).
SNat l
-> (TfoldTree b @@ l)
-> (TfoldTree b @@ l)
-> TfoldTree b @@ (l + 1))
-> RTree d a
-> TfoldTree b @@ d
forall (p :: TyFun Nat Type -> Type) (k :: Nat) a.
KnownNat k =>
Proxy p
-> (a -> p @@ 0)
-> (forall (l :: Nat).
SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1))
-> RTree k a
-> p @@ k
tdfold (Proxy (TfoldTree b)
forall k (t :: k). Proxy t
Proxy @(TfoldTree b)) a -> b
a -> TfoldTree b @@ 0
f ((b -> b -> b) -> SNat l -> b -> b -> b
forall a b. a -> b -> a
const b -> b -> b
g)
treplicate :: forall d a . SNat d -> a -> RTree d a
treplicate :: SNat d -> a -> RTree d a
treplicate sn :: SNat d
sn a :: a
a = UNat d -> RTree d a
forall (n :: Nat). UNat n -> RTree n a
go (SNat d -> UNat d
forall (n :: Nat). SNat n -> UNat n
toUNat SNat d
sn)
where
go :: UNat n -> RTree n a
go :: UNat n -> RTree n a
go UZero = a -> RTree 0 a
forall a. a -> RTree 0 a
LR a
a
go (USucc un :: UNat n
un) = RTree n a -> RTree n a -> RTree (n + 1) a
forall (d :: Nat) a. RTree d a -> RTree d a -> RTree (d + 1) a
BR (UNat n -> RTree n a
forall (n :: Nat). UNat n -> RTree n a
go UNat n
un) (UNat n -> RTree n a
forall (n :: Nat). UNat n -> RTree n a
go UNat n
un)
{-# NOINLINE treplicate #-}
trepeat :: KnownNat d => a -> RTree d a
trepeat :: a -> RTree d a
trepeat = SNat d -> a -> RTree d a
forall (d :: Nat) a. SNat d -> a -> RTree d a
treplicate SNat d
forall (n :: Nat). KnownNat n => SNat n
SNat
data MapTree (a :: Type) (f :: TyFun Nat Type) :: Type
type instance Apply (MapTree a) d = RTree d a
tmap :: forall d a b . KnownNat d => (a -> b) -> RTree d a -> RTree d b
tmap :: (a -> b) -> RTree d a -> RTree d b
tmap f :: a -> b
f = Proxy (MapTree b)
-> (a -> MapTree b @@ 0)
-> (forall (l :: Nat).
SNat l
-> (MapTree b @@ l) -> (MapTree b @@ l) -> MapTree b @@ (l + 1))
-> RTree d a
-> MapTree b @@ d
forall (p :: TyFun Nat Type -> Type) (k :: Nat) a.
KnownNat k =>
Proxy p
-> (a -> p @@ 0)
-> (forall (l :: Nat).
SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1))
-> RTree k a
-> p @@ k
tdfold (Proxy (MapTree b)
forall k (t :: k). Proxy t
Proxy @(MapTree b)) (b -> RTree 0 b
forall a. a -> RTree 0 a
LR (b -> RTree 0 b) -> (a -> b) -> a -> RTree 0 b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f) (\_ l :: MapTree b @@ l
l r :: MapTree b @@ l
r -> RTree l b -> RTree l b -> RTree (l + 1) b
forall (d :: Nat) a. RTree d a -> RTree d a -> RTree (d + 1) a
BR MapTree b @@ l
RTree l b
l MapTree b @@ l
RTree l b
r)
tindices :: forall d . KnownNat d => RTree d (Index (2^d))
tindices :: RTree d (Index (2 ^ d))
tindices =
Proxy (MapTree (Index (2 ^ d)))
-> (Index (2 ^ d) -> MapTree (Index (2 ^ d)) @@ 0)
-> (forall (l :: Nat).
SNat l
-> (MapTree (Index (2 ^ d)) @@ l)
-> (MapTree (Index (2 ^ d)) @@ l)
-> MapTree (Index (2 ^ d)) @@ (l + 1))
-> RTree d (Index (2 ^ d))
-> MapTree (Index (2 ^ d)) @@ d
forall (p :: TyFun Nat Type -> Type) (k :: Nat) a.
KnownNat k =>
Proxy p
-> (a -> p @@ 0)
-> (forall (l :: Nat).
SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1))
-> RTree k a
-> p @@ k
tdfold (Proxy (MapTree (Index (2 ^ d)))
forall k (t :: k). Proxy t
Proxy @(MapTree (Index (2^d)))) Index (2 ^ d) -> MapTree (Index (2 ^ d)) @@ 0
forall a. a -> RTree 0 a
LR
(\s :: SNat l
s@SNat l
SNat l :: MapTree (Index (2 ^ d)) @@ l
l r :: MapTree (Index (2 ^ d)) @@ l
r -> RTree l (Index (2 ^ d))
-> RTree l (Index (2 ^ d)) -> RTree (l + 1) (Index (2 ^ d))
forall (d :: Nat) a. RTree d a -> RTree d a -> RTree (d + 1) a
BR MapTree (Index (2 ^ d)) @@ l
RTree l (Index (2 ^ d))
l ((Index (2 ^ d) -> Index (2 ^ d))
-> RTree l (Index (2 ^ d)) -> RTree l (Index (2 ^ d))
forall (d :: Nat) a b.
KnownNat d =>
(a -> b) -> RTree d a -> RTree d b
tmap (Index (2 ^ d) -> Index (2 ^ d) -> Index (2 ^ d)
forall a. Num a => a -> a -> a
+(SNat (2 ^ l) -> Index (2 ^ d)
forall a (n :: Nat). Num a => SNat n -> a
snatToNum (SNat l -> SNat (2 ^ l)
forall (a :: Nat). SNat a -> SNat (2 ^ a)
pow2SNat SNat l
s))) MapTree (Index (2 ^ d)) @@ l
RTree l (Index (2 ^ d))
r))
(SNat d -> Index (2 ^ d) -> RTree d (Index (2 ^ d))
forall (d :: Nat) a. SNat d -> a -> RTree d a
treplicate SNat d
forall (n :: Nat). KnownNat n => SNat n
SNat 0)
data V2TTree (a :: Type) (f :: TyFun Nat Type) :: Type
type instance Apply (V2TTree a) d = RTree d a
v2t :: forall d a . KnownNat d => Vec (2^d) a -> RTree d a
v2t :: Vec (2 ^ d) a -> RTree d a
v2t = Proxy (V2TTree a)
-> (a -> V2TTree a @@ 0)
-> (forall (l :: Nat).
SNat l
-> (V2TTree a @@ l) -> (V2TTree a @@ l) -> V2TTree a @@ (l + 1))
-> Vec (2 ^ d) a
-> V2TTree a @@ d
forall (p :: TyFun Nat Type -> Type) (k :: Nat) a.
KnownNat k =>
Proxy p
-> (a -> p @@ 0)
-> (forall (l :: Nat).
SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1))
-> Vec (2 ^ k) a
-> p @@ k
dtfold (Proxy (V2TTree a)
forall k (t :: k). Proxy t
Proxy @(V2TTree a)) a -> V2TTree a @@ 0
forall a. a -> RTree 0 a
LR ((RTree l a -> RTree l a -> RTree (l + 1) a)
-> SNat l -> RTree l a -> RTree l a -> RTree (l + 1) a
forall a b. a -> b -> a
const RTree l a -> RTree l a -> RTree (l + 1) a
forall (d :: Nat) a. RTree d a -> RTree d a -> RTree (d + 1) a
BR)
data T2VTree (a :: Type) (f :: TyFun Nat Type) :: Type
type instance Apply (T2VTree a) d = Vec (2^d) a
t2v :: forall d a . KnownNat d => RTree d a -> Vec (2^d) a
t2v :: RTree d a -> Vec (2 ^ d) a
t2v = Proxy (T2VTree a)
-> (a -> T2VTree a @@ 0)
-> (forall (l :: Nat).
SNat l
-> (T2VTree a @@ l) -> (T2VTree a @@ l) -> T2VTree a @@ (l + 1))
-> RTree d a
-> T2VTree a @@ d
forall (p :: TyFun Nat Type -> Type) (k :: Nat) a.
KnownNat k =>
Proxy p
-> (a -> p @@ 0)
-> (forall (l :: Nat).
SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1))
-> RTree k a
-> p @@ k
tdfold (Proxy (T2VTree a)
forall k (t :: k). Proxy t
Proxy @(T2VTree a)) (a -> Vec 0 a -> Vec (0 + 1) a
forall a (n :: Nat). a -> Vec n a -> Vec (n + 1) a
:> Vec 0 a
forall a. Vec 0 a
Nil) (\_ l :: T2VTree a @@ l
l r :: T2VTree a @@ l
r -> T2VTree a @@ l
Vec (2 ^ l) a
l Vec (2 ^ l) a -> Vec (2 ^ l) a -> Vec ((2 ^ l) + (2 ^ l)) a
forall (n :: Nat) a (m :: Nat). Vec n a -> Vec m a -> Vec (n + m) a
++ T2VTree a @@ l
Vec (2 ^ l) a
r)
indexTree :: (KnownNat d, Enum i) => RTree d a -> i -> a
indexTree :: RTree d a -> i -> a
indexTree t :: RTree d a
t i :: i
i = (RTree d a -> Vec (2 ^ d) a
forall (d :: Nat) a. KnownNat d => RTree d a -> Vec (2 ^ d) a
t2v RTree d a
t) Vec (2 ^ d) a -> i -> a
forall (n :: Nat) i a. (KnownNat n, Enum i) => Vec n a -> i -> a
!! i
i
replaceTree :: (KnownNat d, Enum i) => i -> a -> RTree d a -> RTree d a
replaceTree :: i -> a -> RTree d a -> RTree d a
replaceTree i :: i
i a :: a
a = Vec (2 ^ d) a -> RTree d a
forall (d :: Nat) a. KnownNat d => Vec (2 ^ d) a -> RTree d a
v2t (Vec (2 ^ d) a -> RTree d a)
-> (RTree d a -> Vec (2 ^ d) a) -> RTree d a -> RTree d a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> a -> Vec (2 ^ d) a -> Vec (2 ^ d) a
forall (n :: Nat) i a.
(KnownNat n, Enum i) =>
i -> a -> Vec n a -> Vec n a
replace i
i a
a (Vec (2 ^ d) a -> Vec (2 ^ d) a)
-> (RTree d a -> Vec (2 ^ d) a) -> RTree d a -> Vec (2 ^ d) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RTree d a -> Vec (2 ^ d) a
forall (d :: Nat) a. KnownNat d => RTree d a -> Vec (2 ^ d) a
t2v
data ZipWithTree (b :: Type) (c :: Type) (f :: TyFun Nat Type) :: Type
type instance Apply (ZipWithTree b c) d = RTree d b -> RTree d c
tzipWith :: forall a b c d . KnownNat d => (a -> b -> c) -> RTree d a -> RTree d b -> RTree d c
tzipWith :: (a -> b -> c) -> RTree d a -> RTree d b -> RTree d c
tzipWith f :: a -> b -> c
f = Proxy (ZipWithTree b c)
-> (a -> ZipWithTree b c @@ 0)
-> (forall (l :: Nat).
SNat l
-> (ZipWithTree b c @@ l)
-> (ZipWithTree b c @@ l)
-> ZipWithTree b c @@ (l + 1))
-> RTree d a
-> ZipWithTree b c @@ d
forall (p :: TyFun Nat Type -> Type) (k :: Nat) a.
KnownNat k =>
Proxy p
-> (a -> p @@ 0)
-> (forall (l :: Nat).
SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1))
-> RTree k a
-> p @@ k
tdfold (Proxy (ZipWithTree b c)
forall k (t :: k). Proxy t
Proxy @(ZipWithTree b c)) a -> ZipWithTree b c @@ 0
a -> RTree 0 b -> RTree 0 c
lr forall (l :: Nat).
SNat l
-> (ZipWithTree b c @@ l)
-> (ZipWithTree b c @@ l)
-> ZipWithTree b c @@ (l + 1)
forall (l :: Nat).
SNat l
-> (RTree l b -> RTree l c)
-> (RTree l b -> RTree l c)
-> RTree (l + 1) b
-> RTree (l + 1) c
br
where
lr :: a -> RTree 0 b -> RTree 0 c
lr :: a -> RTree 0 b -> RTree 0 c
lr a :: a
a t :: RTree 0 b
t = c -> RTree 0 c
forall a. a -> RTree 0 a
LR (a -> b -> c
f a
a (RTree 0 b -> b
forall a. RTree 0 a -> a
textract RTree 0 b
t))
br :: SNat l
-> (RTree l b -> RTree l c)
-> (RTree l b -> RTree l c)
-> RTree (l+1) b
-> RTree (l+1) c
br :: SNat l
-> (RTree l b -> RTree l c)
-> (RTree l b -> RTree l c)
-> RTree (l + 1) b
-> RTree (l + 1) c
br _ fl :: RTree l b -> RTree l c
fl fr :: RTree l b -> RTree l c
fr t :: RTree (l + 1) b
t = RTree l c -> RTree l c -> RTree (l + 1) c
forall (d :: Nat) a. RTree d a -> RTree d a -> RTree (d + 1) a
BR (RTree l b -> RTree l c
fl RTree l b
l) (RTree l b -> RTree l c
fr RTree l b
r)
where
(l :: RTree l b
l,r :: RTree l b
r) = RTree (l + 1) b -> (RTree l b, RTree l b)
forall (d :: Nat) a. RTree (d + 1) a -> (RTree d a, RTree d a)
tsplit RTree (l + 1) b
t
tzip :: KnownNat d => RTree d a -> RTree d b -> RTree d (a,b)
tzip :: RTree d a -> RTree d b -> RTree d (a, b)
tzip = (a -> b -> (a, b)) -> RTree d a -> RTree d b -> RTree d (a, b)
forall a b c (d :: Nat).
KnownNat d =>
(a -> b -> c) -> RTree d a -> RTree d b -> RTree d c
tzipWith (,)
data UnzipTree (a :: Type) (b :: Type) (f :: TyFun Nat Type) :: Type
type instance Apply (UnzipTree a b) d = (RTree d a, RTree d b)
tunzip :: forall d a b . KnownNat d => RTree d (a,b) -> (RTree d a,RTree d b)
tunzip :: RTree d (a, b) -> (RTree d a, RTree d b)
tunzip = Proxy (UnzipTree a b)
-> ((a, b) -> UnzipTree a b @@ 0)
-> (forall (l :: Nat).
SNat l
-> (UnzipTree a b @@ l)
-> (UnzipTree a b @@ l)
-> UnzipTree a b @@ (l + 1))
-> RTree d (a, b)
-> UnzipTree a b @@ d
forall (p :: TyFun Nat Type -> Type) (k :: Nat) a.
KnownNat k =>
Proxy p
-> (a -> p @@ 0)
-> (forall (l :: Nat).
SNat l -> (p @@ l) -> (p @@ l) -> p @@ (l + 1))
-> RTree k a
-> p @@ k
tdfold (Proxy (UnzipTree a b)
forall k (t :: k). Proxy t
Proxy @(UnzipTree a b)) (a, b) -> UnzipTree a b @@ 0
forall a a. (a, a) -> (RTree 0 a, RTree 0 a)
lr forall p (d :: Nat) a (d :: Nat) a.
p
-> (RTree d a, RTree d a)
-> (RTree d a, RTree d a)
-> (RTree (d + 1) a, RTree (d + 1) a)
forall (l :: Nat).
SNat l
-> (UnzipTree a b @@ l)
-> (UnzipTree a b @@ l)
-> UnzipTree a b @@ (l + 1)
br
where
lr :: (a, a) -> (RTree 0 a, RTree 0 a)
lr (a :: a
a,b :: a
b) = (a -> RTree 0 a
forall a. a -> RTree 0 a
LR a
a,a -> RTree 0 a
forall a. a -> RTree 0 a
LR a
b)
br :: p
-> (RTree d a, RTree d a)
-> (RTree d a, RTree d a)
-> (RTree (d + 1) a, RTree (d + 1) a)
br _ (l1 :: RTree d a
l1,r1 :: RTree d a
r1) (l2 :: RTree d a
l2,r2 :: RTree d a
r2) = (RTree d a -> RTree d a -> RTree (d + 1) a
forall (d :: Nat) a. RTree d a -> RTree d a -> RTree (d + 1) a
BR RTree d a
l1 RTree d a
l2, RTree d a -> RTree d a -> RTree (d + 1) a
forall (d :: Nat) a. RTree d a -> RTree d a -> RTree (d + 1) a
BR RTree d a
r1 RTree d a
r2)
lazyT :: KnownNat d
=> RTree d a
-> RTree d a
lazyT :: RTree d a -> RTree d a
lazyT = (() -> a -> a) -> RTree d () -> RTree d a -> RTree d a
forall a b c (d :: Nat).
KnownNat d =>
(a -> b -> c) -> RTree d a -> RTree d b -> RTree d c
tzipWith ((a -> () -> a) -> () -> a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> () -> a
forall a b. a -> b -> a
const) (() -> RTree d ()
forall (d :: Nat) a. KnownNat d => a -> RTree d a
trepeat ())