{-# Language TupleSections #-}
module Csound.Dynamic.Tfm.UnfoldMultiOuts(
    unfoldMultiOuts, UnfoldMultiOuts(..), Selector(..)
) where

import Data.List(sortBy)
import Data.Ord(comparing)
import Data.Maybe(mapMaybe, isNothing)
import Control.Monad.Trans.State.Strict
import qualified Data.IntMap as IM

import Csound.Dynamic.Tfm.DeduceTypes(Var(..))

type ChildrenMap = IM.IntMap [Port]

lookupChildren :: ChildrenMap -> Var a -> [Port]
lookupChildren :: ChildrenMap -> Var a -> [Port]
lookupChildren ChildrenMap
m Var a
parentVar = ChildrenMap
m ChildrenMap -> Key -> [Port]
forall a. IntMap a -> Key -> a
IM.! Var a -> Key
forall a. Var a -> Key
varId Var a
parentVar

mkChildrenMap :: [(Var a, Selector a)] -> ChildrenMap
mkChildrenMap :: [(Var a, Selector a)] -> ChildrenMap
mkChildrenMap = ([Port] -> [Port] -> [Port]) -> [(Key, [Port])] -> ChildrenMap
forall a. (a -> a -> a) -> [(Key, a)] -> IntMap a
IM.fromListWith [Port] -> [Port] -> [Port]
forall a. [a] -> [a] -> [a]
(++) ([(Key, [Port])] -> ChildrenMap)
-> ([(Var a, Selector a)] -> [(Key, [Port])])
-> [(Var a, Selector a)]
-> ChildrenMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Var a, Selector a) -> (Key, [Port]))
-> [(Var a, Selector a)] -> [(Key, [Port])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Var a, Selector a) -> (Key, [Port])
forall (m :: * -> *) a a.
Monad m =>
(Var a, Selector a) -> (Key, m Port)
extract 
    where extract :: (Var a, Selector a) -> (Key, m Port)
extract (Var a
var, Selector a
sel) = (Var a -> Key
forall a. Var a -> Key
varId (Var a -> Key) -> Var a -> Key
forall a b. (a -> b) -> a -> b
$ Selector a -> Var a
forall a. Selector a -> Var a
selectorParent Selector a
sel, 
                                Port -> m Port
forall (m :: * -> *) a. Monad m => a -> m a
return (Port -> m Port) -> Port -> m Port
forall a b. (a -> b) -> a -> b
$ Key -> Key -> Port
Port (Var a -> Key
forall a. Var a -> Key
varId Var a
var) (Selector a -> Key
forall a. Selector a -> Key
selectorOrder Selector a
sel))

data Port = Port 
    { Port -> Key
portId    :: Int
    , Port -> Key
portOrder :: Int } deriving (Key -> Port -> ShowS
[Port] -> ShowS
Port -> String
(Key -> Port -> ShowS)
-> (Port -> String) -> ([Port] -> ShowS) -> Show Port
forall a.
(Key -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Port] -> ShowS
$cshowList :: [Port] -> ShowS
show :: Port -> String
$cshow :: Port -> String
showsPrec :: Key -> Port -> ShowS
$cshowsPrec :: Key -> Port -> ShowS
Show)

type SingleStmt f a = (Var a, f (Var a))
type MultiStmt  f a = ([Var a], f (Var a))

data Selector a = Selector 
    { Selector a -> Var a
selectorParent  :: Var a
    , Selector a -> Key
selectorOrder   :: Int }

data UnfoldMultiOuts f a = UnfoldMultiOuts {
    UnfoldMultiOuts f a -> f (Var a) -> Maybe (Selector a)
getSelector    :: f (Var a) -> Maybe (Selector a),
    UnfoldMultiOuts f a -> f (Var a) -> Maybe [a]
getParentTypes :: f (Var a) -> Maybe [a] }

unfoldMultiOuts :: UnfoldMultiOuts f a -> Int -> [SingleStmt f a] -> ([MultiStmt f a], Int)
unfoldMultiOuts :: UnfoldMultiOuts f a
-> Key -> [SingleStmt f a] -> ([MultiStmt f a], Key)
unfoldMultiOuts UnfoldMultiOuts f a
algSpec Key
lastFreshId [SingleStmt f a]
stmts = State Key [MultiStmt f a] -> Key -> ([MultiStmt f a], Key)
forall s a. State s a -> s -> (a, s)
runState State Key [MultiStmt f a]
st Key
lastFreshId
    where selectors :: [(Var a, Selector a)]
selectors = (SingleStmt f a -> Maybe (Var a, Selector a))
-> [SingleStmt f a] -> [(Var a, Selector a)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(Var a
lhs, f (Var a)
rhs) -> (Selector a -> (Var a, Selector a))
-> Maybe (Selector a) -> Maybe (Var a, Selector a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Var a
lhs,) (Maybe (Selector a) -> Maybe (Var a, Selector a))
-> Maybe (Selector a) -> Maybe (Var a, Selector a)
forall a b. (a -> b) -> a -> b
$ UnfoldMultiOuts f a -> f (Var a) -> Maybe (Selector a)
forall (f :: * -> *) a.
UnfoldMultiOuts f a -> f (Var a) -> Maybe (Selector a)
getSelector UnfoldMultiOuts f a
algSpec f (Var a)
rhs) [SingleStmt f a]
stmts
          st :: State Key [MultiStmt f a]
st = (SingleStmt f a -> StateT Key Identity (MultiStmt f a))
-> [SingleStmt f a] -> State Key [MultiStmt f a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (UnfoldMultiOuts f a
-> ChildrenMap
-> SingleStmt f a
-> StateT Key Identity (MultiStmt f a)
forall (f :: * -> *) a.
UnfoldMultiOuts f a
-> ChildrenMap -> SingleStmt f a -> State Key (MultiStmt f a)
unfoldStmt UnfoldMultiOuts f a
algSpec (ChildrenMap
 -> SingleStmt f a -> StateT Key Identity (MultiStmt f a))
-> ChildrenMap
-> SingleStmt f a
-> StateT Key Identity (MultiStmt f a)
forall a b. (a -> b) -> a -> b
$ [(Var a, Selector a)] -> ChildrenMap
forall a. [(Var a, Selector a)] -> ChildrenMap
mkChildrenMap [(Var a, Selector a)]
selectors) ([SingleStmt f a] -> State Key [MultiStmt f a])
-> [SingleStmt f a] -> State Key [MultiStmt f a]
forall a b. (a -> b) -> a -> b
$ [SingleStmt f a] -> [SingleStmt f a]
forall a. [(a, f (Var a))] -> [(a, f (Var a))]
dropSelectors [SingleStmt f a]
stmts
          dropSelectors :: [(a, f (Var a))] -> [(a, f (Var a))]
dropSelectors = ((a, f (Var a)) -> Bool) -> [(a, f (Var a))] -> [(a, f (Var a))]
forall a. (a -> Bool) -> [a] -> [a]
filter (Maybe (Selector a) -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe (Selector a) -> Bool)
-> ((a, f (Var a)) -> Maybe (Selector a)) -> (a, f (Var a)) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnfoldMultiOuts f a -> f (Var a) -> Maybe (Selector a)
forall (f :: * -> *) a.
UnfoldMultiOuts f a -> f (Var a) -> Maybe (Selector a)
getSelector UnfoldMultiOuts f a
algSpec (f (Var a) -> Maybe (Selector a))
-> ((a, f (Var a)) -> f (Var a))
-> (a, f (Var a))
-> Maybe (Selector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, f (Var a)) -> f (Var a)
forall a b. (a, b) -> b
snd)

unfoldStmt :: UnfoldMultiOuts f a -> ChildrenMap -> SingleStmt f a -> State Int (MultiStmt f a)
unfoldStmt :: UnfoldMultiOuts f a
-> ChildrenMap -> SingleStmt f a -> State Key (MultiStmt f a)
unfoldStmt UnfoldMultiOuts f a
algSpec ChildrenMap
childrenMap (Var a
lhs, f (Var a)
rhs) = case UnfoldMultiOuts f a -> f (Var a) -> Maybe [a]
forall (f :: * -> *) a.
UnfoldMultiOuts f a -> f (Var a) -> Maybe [a]
getParentTypes UnfoldMultiOuts f a
algSpec f (Var a)
rhs of
    Maybe [a]
Nothing    -> MultiStmt f a -> State Key (MultiStmt f a)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Var a
lhs], f (Var a)
rhs)
    Just [a]
types -> ([Var a] -> MultiStmt f a)
-> StateT Key Identity [Var a] -> State Key (MultiStmt f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (,f (Var a)
rhs) (StateT Key Identity [Var a] -> State Key (MultiStmt f a))
-> StateT Key Identity [Var a] -> State Key (MultiStmt f a)
forall a b. (a -> b) -> a -> b
$ [Port] -> [a] -> StateT Key Identity [Var a]
forall a. [Port] -> [a] -> State Key [Var a]
formLhs (ChildrenMap -> Var a -> [Port]
forall a. ChildrenMap -> Var a -> [Port]
lookupChildren ChildrenMap
childrenMap Var a
lhs) [a]
types

formLhs :: [Port] -> [a] -> State Int [Var a]
formLhs :: [Port] -> [a] -> State Key [Var a]
formLhs [Port]
ports [a]
types = ([Key] -> [Var a])
-> StateT Key Identity [Key] -> State Key [Var a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> Key -> Var a) -> [a] -> [Key] -> [Var a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((Key -> a -> Var a) -> a -> Key -> Var a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Key -> a -> Var a
forall a. Key -> a -> Var a
Var) [a]
types) ([Port] -> StateT Key Identity [Key]
forall (m :: * -> *). Monad m => [Port] -> StateT Key m [Key]
getPorts [Port]
ports)
    where getPorts :: [Port] -> StateT Key m [Key]
getPorts [Port]
ps = (Key -> ([Key], Key)) -> StateT Key m [Key]
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((Key -> ([Key], Key)) -> StateT Key m [Key])
-> (Key -> ([Key], Key)) -> StateT Key m [Key]
forall a b. (a -> b) -> a -> b
$ \Key
lastFreshId -> 
            let ps' :: [Port]
ps' = (Port -> Port -> Ordering) -> [Port] -> [Port]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy ((Port -> Key) -> Port -> Port -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing Port -> Key
portOrder) [Port]
ps
                ([[Key]]
ids, Key
lastPortOrder) = State Key [[Key]] -> Key -> ([[Key]], Key)
forall s a. State s a -> s -> (a, s)
runState ((Port -> StateT Key Identity [Key]) -> [Port] -> State Key [[Key]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Key -> Port -> StateT Key Identity [Key]
fillMissingPorts Key
lastFreshId) [Port]
ps') Key
0
                freshIdForTail :: Key
freshIdForTail = Key
1 Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
lastFreshId Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
inUsePortsSize
                tailIds :: [Key]
tailIds = (Key -> Key) -> [Key] -> [Key]
forall a b. (a -> b) -> [a] -> [b]
map (Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
freshIdForTail) [Key
0 .. Key
outputArity Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1 Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
lastPortOrder]
            in  ([[Key]] -> [Key]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Key]]
ids [Key] -> [Key] -> [Key]
forall a. [a] -> [a] -> [a]
++ [Key]
tailIds, Key
lastFreshId Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
outputArity Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
inUsePortsSize)

          outputArity :: Key
outputArity = [a] -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length [a]
types    
          inUsePortsSize :: Key
inUsePortsSize = [Port] -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length [Port]
ports  
                                
          fillMissingPorts :: Int -> Port -> State Int [Int]
          fillMissingPorts :: Key -> Port -> StateT Key Identity [Key]
fillMissingPorts Key
lastFreshId Port
port = (Key -> ([Key], Key)) -> StateT Key Identity [Key]
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((Key -> ([Key], Key)) -> StateT Key Identity [Key])
-> (Key -> ([Key], Key)) -> StateT Key Identity [Key]
forall a b. (a -> b) -> a -> b
$ \Key
s ->
                if Key
s Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
order
                then ([Key
e], Key
next) 
                else ((Key -> Key) -> [Key] -> [Key]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
lastFreshId) [Key
s .. Key
order Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1] [Key] -> [Key] -> [Key]
forall a. [a] -> [a] -> [a]
++ [Key
e], Key
next)
            where e :: Key
e = Port -> Key
portId Port
port
                  order :: Key
order = Port -> Key
portOrder Port
port                  
                  next :: Key
next = Key
order Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
1