{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TupleSections #-}

module Language.Fortran.Vars.Memory
  ( allocateMemoryBlocks
  , processCommon
  , getTypeSize
  )
where

import           Language.Fortran.Extras        ( allPUS )
import           Data.Data                      ( Data )
import           Data.List                      ( foldl' )
import           Data.Maybe                     ( fromMaybe )
import qualified Data.Map                      as M
import           Language.Fortran.Analysis      ( Analysis
                                                , srcName
                                                )
import           Language.Fortran.AST           ( AList(..)
                                                , aStrip
                                                , CommonGroup(..)
                                                , Expression(..)
                                                , Name
                                                , ProgramUnit
                                                , Statement(..)
                                                , Declarator(..)
                                                )

import           Language.Fortran.Vars.MemoryLocation
                                                ( getStartLocation )
import           Language.Fortran.Vars.Types    ( SymbolTableEntry(..)
                                                , MemoryBlock(..)
                                                , ProgramUnitModel
                                                , SymbolTable
                                                , StorageClass(..)
                                                , StorageTable
                                                , Type
                                                , SemType(..)
                                                )
import           Language.Fortran.Vars.Kind     ( getTypeKind )
import           Language.Fortran.Vars.Union    ( union )

-- | Given a 'SymbolTable' and an 'Expression', return the size of
-- the variable represented by the expression
getSize :: Data a => SymbolTable -> Expression (Analysis a) -> Int
getSize :: forall a. Data a => SymbolTable -> Expression (Analysis a) -> Int
getSize SymbolTable
symTable Expression (Analysis a)
expr =
  let symbol :: Name
symbol = case Expression (Analysis a)
expr of
        ExpValue{} -> Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
expr
        ExpFunctionCall Analysis a
_ SrcSpan
_ e :: Expression (Analysis a)
e@ExpValue{} AList Argument (Analysis a)
_ -> Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e
        ExpSubscript Analysis a
_ SrcSpan
_ e :: Expression (Analysis a)
e@ExpValue{} AList Index (Analysis a)
_ -> Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e
        Expression (Analysis a)
_          -> Name -> Name
forall a. HasCallStack => Name -> a
error Name
"Unsupported expression"
      Just SymbolTableEntry
entity = Name -> SymbolTable -> Maybe SymbolTableEntry
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
symbol SymbolTable
symTable
  in  case SymbolTableEntry
entity of
        SVariable Type
ty Location
_ -> Type -> Int
getTypeSize Type
ty
        SymbolTableEntry
_              -> Name -> Int
forall a. HasCallStack => Name -> a
error (Name
symbol Name -> Name -> Name
forall a. [a] -> [a] -> [a]
++ Name
" is not a VariableEntry.")

getTypeSize :: Type -> Int
getTypeSize :: Type -> Int
getTypeSize = \case
  TArray Type
ty Maybe Dimensions
dims ->
    Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Name -> Int
forall a. HasCallStack => Name -> a
error Name
"Can't calculate size of dynamic array")
      (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$   Int -> Dimensions -> Int
sizeOfStaticArray
      (Int -> Dimensions -> Int)
-> Maybe Int -> Maybe (Dimensions -> Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Maybe Int
getTypeKind Type
ty
      Maybe (Dimensions -> Int) -> Maybe Dimensions -> Maybe Int
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Dimensions
dims
  Type
ty -> Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Name -> Int
forall a. HasCallStack => Name -> a
error Name
"Can't get size of dynamic variable") (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Type -> Maybe Int
getTypeKind Type
ty

-- | Given a static array's 'kind' and 'dimension', calculate its size
sizeOfStaticArray :: Int -> [(Int, Int)] -> Int
sizeOfStaticArray :: Int -> Dimensions -> Int
sizeOfStaticArray Int
kind' Dimensions
dimension' =
  let arraySize :: Int
arraySize = (Int -> (Int, Int) -> Int) -> Int -> Dimensions -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Int
acc (Int
l, Int
h) -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Int
1 Dimensions
dimension'
  in  Int
kind' Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
arraySize

-- | Given a 'SymbolTable', generate a 'StorageTable' for the 'SymbolTable' where
-- each symbol has been assinged to a 'MemoryBlock' within the 'StorageTable' so
-- long as it is not constant
allocateMemoryBlocks :: SymbolTable -> StorageTable
allocateMemoryBlocks :: SymbolTable -> StorageTable
allocateMemoryBlocks = (StorageTable -> Name -> SymbolTableEntry -> StorageTable)
-> StorageTable -> SymbolTable -> StorageTable
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey StorageTable -> Name -> SymbolTableEntry -> StorageTable
f StorageTable
forall k a. Map k a
M.empty
 where
  f :: StorageTable -> Name -> SymbolTableEntry -> StorageTable
  f :: StorageTable -> Name -> SymbolTableEntry -> StorageTable
f StorageTable
storageTable Name
symbol SymbolTableEntry
entry = case SymbolTableEntry
entry of
    SVariable Type
ty Location
_ ->
      let size :: Maybe Int
size = case Type
ty of
            TArray Type
ty' Maybe Dimensions
dims -> Int -> Dimensions -> Int
sizeOfStaticArray (Int -> Dimensions -> Int)
-> Maybe Int -> Maybe (Dimensions -> Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> Maybe Int
getTypeKind Type
ty' Maybe (Dimensions -> Int) -> Maybe Dimensions -> Maybe Int
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Dimensions
dims
            Type
_               -> Type -> Maybe Int
getTypeKind Type
ty
          block :: MemoryBlock
block = MemoryBlock :: Maybe Int -> StorageClass -> [Name] -> MemoryBlock
MemoryBlock
            { blockSize :: Maybe Int
blockSize    = Maybe Int
size
            , storageClass :: StorageClass
storageClass = case Maybe Int
size of
                               Maybe Int
Nothing -> StorageClass
Automatic
                               Maybe Int
_       -> StorageClass
Unspecified
            , variables :: [Name]
variables    = [Name
symbol]
            }
      in  Name -> MemoryBlock -> StorageTable -> StorageTable
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
symbol MemoryBlock
block StorageTable
storageTable
    SymbolTableEntry
_ -> StorageTable
storageTable

-- | Given a 'ProgramUnit' and a 'ProgramUnitModel', resolve any commonly defined global
-- variables in the 'ProgramUnit's to be the same in both memory and within the 'SymbolTable'
processCommon
  :: Data a => ProgramUnit (Analysis a) -> ProgramUnitModel -> ProgramUnitModel
processCommon :: forall a.
Data a =>
ProgramUnit (Analysis a) -> ProgramUnitModel -> ProgramUnitModel
processCommon ProgramUnit (Analysis a)
pu ProgramUnitModel
puModel =
  let commonGrps :: [[CommonGroup (Analysis a)]]
commonGrps =
          [ [CommonGroup (Analysis a)]
commGrps | (StCommon Analysis a
_ SrcSpan
_ (AList Analysis a
_ SrcSpan
_ [CommonGroup (Analysis a)]
commGrps)) <- ProgramUnit (Analysis a) -> [Statement (Analysis a)]
forall a. Data a => ProgramUnit a -> [Statement a]
allPUS ProgramUnit (Analysis a)
pu ]
      mergeCommonVariables :: Map Name [Declarator (Analysis a)]
-> CommonGroup (Analysis a) -> Map Name [Declarator (Analysis a)]
mergeCommonVariables Map Name [Declarator (Analysis a)]
mapping (CommonGroup Analysis a
_ SrcSpan
_ Maybe (Expression (Analysis a))
commName AList Declarator (Analysis a)
decls) =
          let commonName :: Name
commonName = case Maybe (Expression (Analysis a))
commName of
                Just Expression (Analysis a)
e  -> Name
"/" Name -> Name -> Name
forall a. [a] -> [a] -> [a]
++ Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e Name -> Name -> Name
forall a. [a] -> [a] -> [a]
++ Name
"/"
                Maybe (Expression (Analysis a))
Nothing -> Name
"*blank_common*"
              precedingDecls :: [Declarator (Analysis a)]
precedingDecls = [Declarator (Analysis a)]
-> Maybe [Declarator (Analysis a)] -> [Declarator (Analysis a)]
forall a. a -> Maybe a -> a
fromMaybe [] (Name
-> Map Name [Declarator (Analysis a)]
-> Maybe [Declarator (Analysis a)]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
commonName Map Name [Declarator (Analysis a)]
mapping)
          in  Name
-> [Declarator (Analysis a)]
-> Map Name [Declarator (Analysis a)]
-> Map Name [Declarator (Analysis a)]
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
commonName ([Declarator (Analysis a)]
precedingDecls [Declarator (Analysis a)]
-> [Declarator (Analysis a)] -> [Declarator (Analysis a)]
forall a. [a] -> [a] -> [a]
++ AList Declarator (Analysis a) -> [Declarator (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Declarator (Analysis a)
decls) Map Name [Declarator (Analysis a)]
mapping
      commons :: Map Name [Declarator (Analysis a)]
commons = (Map Name [Declarator (Analysis a)]
 -> CommonGroup (Analysis a) -> Map Name [Declarator (Analysis a)])
-> Map Name [Declarator (Analysis a)]
-> [CommonGroup (Analysis a)]
-> Map Name [Declarator (Analysis a)]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map Name [Declarator (Analysis a)]
-> CommonGroup (Analysis a) -> Map Name [Declarator (Analysis a)]
forall {a}.
Map Name [Declarator (Analysis a)]
-> CommonGroup (Analysis a) -> Map Name [Declarator (Analysis a)]
mergeCommonVariables Map Name [Declarator (Analysis a)]
forall k a. Map k a
M.empty ([[CommonGroup (Analysis a)]] -> [CommonGroup (Analysis a)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[CommonGroup (Analysis a)]]
commonGrps)
      processComm :: Name
-> [Declarator (Analysis a)]
-> ProgramUnitModel
-> ProgramUnitModel
processComm Name
commonName [Declarator (Analysis a)]
varDecls (SymbolTable
symTable, StorageTable
mbs) =
          let varExps :: [Expression (Analysis a)]
varExps            = (Declarator (Analysis a) -> Expression (Analysis a))
-> [Declarator (Analysis a)] -> [Expression (Analysis a)]
forall a b. (a -> b) -> [a] -> [b]
map Declarator (Analysis a) -> Expression (Analysis a)
forall {a}. Declarator a -> Expression a
declExpr [Declarator (Analysis a)]
varDecls
              varLocations :: [Location]
varLocations       = (Expression (Analysis a) -> Location)
-> [Expression (Analysis a)] -> [Location]
forall a b. (a -> b) -> [a] -> [b]
map (SymbolTable -> Expression (Analysis a) -> Location
forall a.
Data a =>
SymbolTable -> Expression (Analysis a) -> Location
getStartLocation SymbolTable
symTable) [Expression (Analysis a)]
varExps
              varSizes :: [Int]
varSizes           = (Expression (Analysis a) -> Int)
-> [Expression (Analysis a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (SymbolTable -> Expression (Analysis a) -> Int
forall a. Data a => SymbolTable -> Expression (Analysis a) -> Int
getSize SymbolTable
symTable) [Expression (Analysis a)]
varExps
              varAccumSizes :: [Int]
varAccumSizes      = (Int -> Int -> Int) -> [Int] -> [Int]
forall a. (a -> a -> a) -> [a] -> [a]
scanl1 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [Int]
varSizes
              commBlockLocations :: [Location]
commBlockLocations = (Int -> Location) -> [Int] -> [Location]
forall a b. (a -> b) -> [a] -> [b]
map (Name
commonName, ) (Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
varAccumSizes)
              mbs' :: StorageTable
mbs'               = case Name -> StorageTable -> Maybe MemoryBlock
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
commonName StorageTable
mbs of
                Just MemoryBlock
_ -> StorageTable
mbs
                Maybe MemoryBlock
Nothing ->
                  let newBlock :: MemoryBlock
newBlock = MemoryBlock :: Maybe Int -> StorageClass -> [Name] -> MemoryBlock
MemoryBlock { blockSize :: Maybe Int
blockSize = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
varSizes
                                             , storageClass :: StorageClass
storageClass = StorageClass
Common
                                             , variables :: [Name]
variables = []
                                             }
                  in  Name -> MemoryBlock -> StorageTable -> StorageTable
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
commonName MemoryBlock
newBlock StorageTable
mbs
              f :: ProgramUnitModel -> (Location, Location) -> ProgramUnitModel
f ProgramUnitModel
model (Location
l1, Location
l2) = let (ProgramUnitModel
model', Location
_) = ProgramUnitModel
-> Location -> Location -> (ProgramUnitModel, Location)
union ProgramUnitModel
model Location
l1 Location
l2 in ProgramUnitModel
model'
          in  (ProgramUnitModel -> (Location, Location) -> ProgramUnitModel)
-> ProgramUnitModel -> [(Location, Location)] -> ProgramUnitModel
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ProgramUnitModel -> (Location, Location) -> ProgramUnitModel
f (SymbolTable
symTable, StorageTable
mbs') ([Location] -> [Location] -> [(Location, Location)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Location]
commBlockLocations [Location]
varLocations)
      declExpr :: Declarator a -> Expression a
declExpr (Declarator a
_ SrcSpan
_ Expression a
e DeclaratorType a
_ Maybe (Expression a)
_ Maybe (Expression a)
_) = Expression a
e
  in  (Name
 -> [Declarator (Analysis a)]
 -> ProgramUnitModel
 -> ProgramUnitModel)
-> ProgramUnitModel
-> Map Name [Declarator (Analysis a)]
-> ProgramUnitModel
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
M.foldrWithKey Name
-> [Declarator (Analysis a)]
-> ProgramUnitModel
-> ProgramUnitModel
forall {a}.
Data a =>
Name
-> [Declarator (Analysis a)]
-> ProgramUnitModel
-> ProgramUnitModel
processComm ProgramUnitModel
puModel Map Name [Declarator (Analysis a)]
commons