-- | Infinite lazy partition tables (used for caching).
--
-- We cache almost all computations (which would be otherwise typically 
-- executed many times); this really helps performance.
--

{-# LANGUAGE Rank2Types #-} 
module Math.RootLoci.Misc.PTable where

--------------------------------------------------------------------------------

import Data.List

import Math.Combinat.Classes
import Math.Combinat.Partitions.Integer
import Math.Combinat.Partitions.Set

import qualified Data.Map.Lazy as LMap

import Math.RootLoci.Algebra.SymmPoly

--------------------------------------------------------------------------------
-- * Finite lazy partition tables

newtype PTable a = PTable (LMap.Map Partition a)

createPTable :: (Partition -> a) -> Int -> PTable a
createPTable :: (Partition -> a) -> Int -> PTable a
createPTable Partition -> a
f Int
n = Map Partition a -> PTable a
forall a. Map Partition a -> PTable a
PTable (Map Partition a -> PTable a) -> Map Partition a -> PTable a
forall a b. (a -> b) -> a -> b
$ [(Partition, a)] -> Map Partition a
forall k a. Ord k => [(k, a)] -> Map k a
LMap.fromList [ (Partition
p, Partition -> a
f Partition
p) | Partition
p <- Int -> [Partition]
partitions Int
n ]

lookupPTable :: Partition -> PTable a -> a
lookupPTable :: Partition -> PTable a -> a
lookupPTable Partition
p (PTable Map Partition a
lmap) = case Partition -> Map Partition a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
LMap.lookup Partition
p Map Partition a
lmap of
  Just a
y  -> a
y
  Maybe a
Nothing -> [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"lookupPTable"

--------------------------------------------------------------------------------
-- * Infinite lazy partition tables

newtype PSeries a = PSeries [PTable a]
  
createPSeries :: (Partition -> a) -> PSeries a
createPSeries :: (Partition -> a) -> PSeries a
createPSeries Partition -> a
f = [PTable a] -> PSeries a
forall a. [PTable a] -> PSeries a
PSeries [ (Partition -> a) -> Int -> PTable a
forall a. (Partition -> a) -> Int -> PTable a
createPTable Partition -> a
f Int
n | Int
n<-[Int
0..] ]

lookupPSeries :: Partition -> PSeries a -> a
lookupPSeries :: Partition -> PSeries a -> a
lookupPSeries Partition
part (PSeries [PTable a]
series) = Partition -> PTable a -> a
forall a. Partition -> PTable a -> a
lookupPTable Partition
part ([PTable a]
series [PTable a] -> Int -> PTable a
forall a. [a] -> Int -> a
!! Partition -> Int
partitionWeight Partition
part)

--------------------------------------------------------------------------------
-- * Finite lazy set-partition tables

newtype SetPTable a = SetPTable (LMap.Map SetPartition a)

createSetPTable :: (SetPartition -> a) -> Int -> SetPTable a
createSetPTable :: (SetPartition -> a) -> Int -> SetPTable a
createSetPTable SetPartition -> a
f Int
n = Map SetPartition a -> SetPTable a
forall a. Map SetPartition a -> SetPTable a
SetPTable (Map SetPartition a -> SetPTable a)
-> Map SetPartition a -> SetPTable a
forall a b. (a -> b) -> a -> b
$ [(SetPartition, a)] -> Map SetPartition a
forall k a. Ord k => [(k, a)] -> Map k a
LMap.fromList [ (SetPartition
p, SetPartition -> a
f SetPartition
p) | SetPartition
p <- Int -> [SetPartition]
setPartitions Int
n ]

lookupSetPTable :: SetPartition -> SetPTable a -> a
lookupSetPTable :: SetPartition -> SetPTable a -> a
lookupSetPTable SetPartition
p (SetPTable Map SetPartition a
lmap) = case SetPartition -> Map SetPartition a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
LMap.lookup SetPartition
p Map SetPartition a
lmap of
  Just a
y  -> a
y
  Maybe a
Nothing -> [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"lookupSetPTable"

--------------------------------------------------------------------------------
-- * Infinite lazy set-partition tables

newtype SetPSeries a = SetPSeries [SetPTable a]
  
createSetPSeries :: (SetPartition -> a) -> SetPSeries a
createSetPSeries :: (SetPartition -> a) -> SetPSeries a
createSetPSeries SetPartition -> a
f = [SetPTable a] -> SetPSeries a
forall a. [SetPTable a] -> SetPSeries a
SetPSeries [ (SetPartition -> a) -> Int -> SetPTable a
forall a. (SetPartition -> a) -> Int -> SetPTable a
createSetPTable SetPartition -> a
f Int
n | Int
n<-[Int
0..] ]

lookupSetPSeries :: SetPartition -> SetPSeries a -> a
lookupSetPSeries :: SetPartition -> SetPSeries a -> a
lookupSetPSeries SetPartition
setp (SetPSeries [SetPTable a]
series) = SetPartition -> SetPTable a -> a
forall a. SetPartition -> SetPTable a -> a
lookupSetPTable SetPartition
setp ([SetPTable a]
series [SetPTable a] -> Int -> SetPTable a
forall a. [a] -> Int -> a
!! SetPartition -> Int
setpWeight SetPartition
setp) where
  setpWeight :: SetPartition -> Int
setpWeight (SetPartition [[Int]]
ps) = (Int -> Int -> Int) -> Int -> [Int] -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 (([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Int]]
ps)

--------------------------------------------------------------------------------
-- * polymorphic caching 

polyCache1 
  :: (CacheKey key) 
  => (forall base. ChernBase base => key -> f base)     -- ^ polymorphic function to be cached
  -> (forall base. ChernBase base => key -> f base)
polyCache1 :: (forall base. ChernBase base => key -> f base)
-> forall base. ChernBase base => key -> f base
polyCache1 forall base. ChernBase base => key -> f base
calc = \key
key -> (f AB, f Chern) -> ChernBase base => f base
forall (f :: * -> *) base.
(f AB, f Chern) -> ChernBase base => f base
select1 (key -> f AB
cacheAB key
key, key -> f Chern
cacheChern key
key)  where
  cacheAB :: key -> f AB
cacheAB    = (key -> f AB) -> key -> f AB
forall key a. CacheKey key => (key -> a) -> key -> a
monoCache ((key -> f AB) -> key -> f AB) -> (key -> f AB) -> key -> f AB
forall a b. (a -> b) -> a -> b
$ \key
k -> Sing AB -> (forall b. ChernBase b => f b) -> f AB
forall base (f :: * -> *).
ChernBase base =>
Sing base -> (forall b. ChernBase b => f b) -> f base
spec1' Sing AB
ChernRoot  (key -> f b
forall base. ChernBase base => key -> f base
calc key
k)
  cacheChern :: key -> f Chern
cacheChern = (key -> f Chern) -> key -> f Chern
forall key a. CacheKey key => (key -> a) -> key -> a
monoCache ((key -> f Chern) -> key -> f Chern)
-> (key -> f Chern) -> key -> f Chern
forall a b. (a -> b) -> a -> b
$ \key
k -> Sing Chern -> (forall b. ChernBase b => f b) -> f Chern
forall base (f :: * -> *).
ChernBase base =>
Sing base -> (forall b. ChernBase b => f b) -> f base
spec1' Sing Chern
ChernClass (key -> f b
forall base. ChernBase base => key -> f base
calc key
k)

polyCache2 
  :: (CacheKey key) 
  => (forall base. ChernBase base => key -> f (g base))     -- ^ polymorphic function to be cached
  -> (forall base. ChernBase base => key -> f (g base))
polyCache2 :: (forall base. ChernBase base => key -> f (g base))
-> forall base. ChernBase base => key -> f (g base)
polyCache2 forall base. ChernBase base => key -> f (g base)
calc = \key
key -> (f (g AB), f (g Chern)) -> ChernBase base => f (g base)
forall (f :: * -> *) (g :: * -> *) base.
(f (g AB), f (g Chern)) -> ChernBase base => f (g base)
select2 (key -> f (g AB)
cacheAB key
key, key -> f (g Chern)
cacheChern key
key)  where
  cacheAB :: key -> f (g AB)
cacheAB    = (key -> f (g AB)) -> key -> f (g AB)
forall key a. CacheKey key => (key -> a) -> key -> a
monoCache ((key -> f (g AB)) -> key -> f (g AB))
-> (key -> f (g AB)) -> key -> f (g AB)
forall a b. (a -> b) -> a -> b
$ \key
k -> Sing AB -> (forall b. ChernBase b => f (g b)) -> f (g AB)
forall base (f :: * -> *) (g :: * -> *).
ChernBase base =>
Sing base -> (forall b. ChernBase b => f (g b)) -> f (g base)
spec2' Sing AB
ChernRoot  (key -> f (g b)
forall base. ChernBase base => key -> f (g base)
calc key
k)
  cacheChern :: key -> f (g Chern)
cacheChern = (key -> f (g Chern)) -> key -> f (g Chern)
forall key a. CacheKey key => (key -> a) -> key -> a
monoCache ((key -> f (g Chern)) -> key -> f (g Chern))
-> (key -> f (g Chern)) -> key -> f (g Chern)
forall a b. (a -> b) -> a -> b
$ \key
k -> Sing Chern -> (forall b. ChernBase b => f (g b)) -> f (g Chern)
forall base (f :: * -> *) (g :: * -> *).
ChernBase base =>
Sing base -> (forall b. ChernBase b => f (g b)) -> f (g base)
spec2' Sing Chern
ChernClass (key -> f (g b)
forall base. ChernBase base => key -> f (g base)
calc key
k)

polyCache3 
  :: (CacheKey key) 
  => (forall base. ChernBase base => key -> f (g (h base)))     -- ^ polymorphic function to be cached
  -> (forall base. ChernBase base => key -> f (g (h base)))
polyCache3 :: (forall base. ChernBase base => key -> f (g (h base)))
-> forall base. ChernBase base => key -> f (g (h base))
polyCache3 forall base. ChernBase base => key -> f (g (h base))
calc = \key
key -> (f (g (h AB)), f (g (h Chern))) -> ChernBase base => f (g (h base))
forall (f :: * -> *) (g :: * -> *) (h :: * -> *) base.
(f (g (h AB)), f (g (h Chern))) -> ChernBase base => f (g (h base))
select3 (key -> f (g (h AB))
cacheAB key
key, key -> f (g (h Chern))
cacheChern key
key)  where
  cacheAB :: key -> f (g (h AB))
cacheAB    = (key -> f (g (h AB))) -> key -> f (g (h AB))
forall key a. CacheKey key => (key -> a) -> key -> a
monoCache ((key -> f (g (h AB))) -> key -> f (g (h AB)))
-> (key -> f (g (h AB))) -> key -> f (g (h AB))
forall a b. (a -> b) -> a -> b
$ \key
k -> Sing AB -> (forall b. ChernBase b => f (g (h b))) -> f (g (h AB))
forall base (f :: * -> *) (g :: * -> *) (h :: * -> *).
ChernBase base =>
Sing base
-> (forall b. ChernBase b => f (g (h b))) -> f (g (h base))
spec3' Sing AB
ChernRoot  (key -> f (g (h b))
forall base. ChernBase base => key -> f (g (h base))
calc key
k)
  cacheChern :: key -> f (g (h Chern))
cacheChern = (key -> f (g (h Chern))) -> key -> f (g (h Chern))
forall key a. CacheKey key => (key -> a) -> key -> a
monoCache ((key -> f (g (h Chern))) -> key -> f (g (h Chern)))
-> (key -> f (g (h Chern))) -> key -> f (g (h Chern))
forall a b. (a -> b) -> a -> b
$ \key
k -> Sing Chern
-> (forall b. ChernBase b => f (g (h b))) -> f (g (h Chern))
forall base (f :: * -> *) (g :: * -> *) (h :: * -> *).
ChernBase base =>
Sing base
-> (forall b. ChernBase b => f (g (h b))) -> f (g (h base))
spec3' Sing Chern
ChernClass (key -> f (g (h b))
forall base. ChernBase base => key -> f (g (h base))
calc key
k)

--------------------------------------------------------------------------------
-- * monomorphic caching 

class CacheKey key where
  monoCache :: (key -> a) -> (key -> a)

instance CacheKey Int          where  monoCache :: (Int -> a) -> Int -> a
monoCache = (Int -> a) -> Int -> a
forall a. (Int -> a) -> Int -> a
icache
instance CacheKey Partition    where  monoCache :: (Partition -> a) -> Partition -> a
monoCache = (Partition -> a) -> Partition -> a
forall a. (Partition -> a) -> Partition -> a
pcache
instance CacheKey SetPartition where  monoCache :: (SetPartition -> a) -> SetPartition -> a
monoCache = (SetPartition -> a) -> SetPartition -> a
forall a. (SetPartition -> a) -> SetPartition -> a
setpcache

--------------------------------------------------------------------------------
-- * individual caching functions

pcache :: (Partition -> a) -> (Partition -> a)
pcache :: (Partition -> a) -> Partition -> a
pcache Partition -> a
calc = Partition -> a
lkp where
  lkp :: Partition -> a
lkp Partition
p = Partition -> PSeries a -> a
forall a. Partition -> PSeries a -> a
lookupPSeries Partition
p PSeries a
table 
  table :: PSeries a
table = (Partition -> a) -> PSeries a
forall a. (Partition -> a) -> PSeries a
createPSeries Partition -> a
calc

setpcache :: (SetPartition -> a) -> (SetPartition -> a)
setpcache :: (SetPartition -> a) -> SetPartition -> a
setpcache SetPartition -> a
calc = SetPartition -> a
lkp where
  lkp :: SetPartition -> a
lkp SetPartition
setp = SetPartition -> SetPSeries a -> a
forall a. SetPartition -> SetPSeries a -> a
lookupSetPSeries SetPartition
setp SetPSeries a
table 
  table :: SetPSeries a
table    = (SetPartition -> a) -> SetPSeries a
forall a. (SetPartition -> a) -> SetPSeries a
createSetPSeries SetPartition -> a
calc

icache :: (Int -> a) -> (Int -> a)
icache :: (Int -> a) -> Int -> a
icache Int -> a
calc = \Int
n -> ([a]
table [a] -> Int -> a
forall a. [a] -> Int -> a
!! Int
n) where
  table :: [a]
table = [ Int -> a
calc Int
i | Int
i <- [Int
0..]  ]

icache' :: a -> Int -> (Int -> a) -> (Int -> a)
icache' :: a -> Int -> (Int -> a) -> Int -> a
icache' a
dflt Int
fstidx Int -> a
calc = \Int
n -> if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
fstidx then a
dflt else ([a]
table [a] -> Int -> a
forall a. [a] -> Int -> a
!! (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
fstidx)) where
  table :: [a]
table = [ Int -> a
calc Int
i | Int
i <- [Int
fstidx..]  ]

--------------------------------------------------------------------------------