{-# 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 )
import           Language.Fortran.Analysis.SemanticTypes
                                                ( dimensionsToTuples )

-- | 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 :: MemoryBlockName
symbol = case Expression (Analysis a)
expr of
        ExpValue{} -> forall a. Expression (Analysis a) -> MemoryBlockName
srcName Expression (Analysis a)
expr
        ExpFunctionCall Analysis a
_ SrcSpan
_ e :: Expression (Analysis a)
e@ExpValue{} AList Argument (Analysis a)
_ -> forall a. Expression (Analysis a) -> MemoryBlockName
srcName Expression (Analysis a)
e
        ExpSubscript Analysis a
_ SrcSpan
_ e :: Expression (Analysis a)
e@ExpValue{} AList Index (Analysis a)
_ -> forall a. Expression (Analysis a) -> MemoryBlockName
srcName Expression (Analysis a)
e
        Expression (Analysis a)
_          -> forall a. HasCallStack => MemoryBlockName -> a
error MemoryBlockName
"Unsupported expression"
      Just SymbolTableEntry
entity = forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup MemoryBlockName
symbol SymbolTable
symTable
  in  case SymbolTableEntry
entity of
        SVariable Type
ty Location
_ -> Type -> Int
getTypeSize Type
ty
        SymbolTableEntry
_              -> forall a. HasCallStack => MemoryBlockName -> a
error (MemoryBlockName
symbol forall a. [a] -> [a] -> [a]
++ MemoryBlockName
" is not a VariableEntry.")

getTypeSize :: Type -> Int
getTypeSize :: Type -> Int
getTypeSize =
    forall a. a -> Maybe a -> a
fromMaybe (forall a. HasCallStack => MemoryBlockName -> a
error MemoryBlockName
"Can't get size of dynamic variable") forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Maybe Int
getTypeSize'

getTypeSize' :: Type -> Maybe Int
getTypeSize' :: Type -> Maybe Int
getTypeSize' = \case
  TArray Type
ty Dimensions
dims -> do
    [(Int, Int)]
dims' <- Dimensions -> Maybe [(Int, Int)]
dimensionsToTuples Dimensions
dims
    Int
kind <- Type -> Maybe Int
getTypeKind Type
ty
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int -> [(Int, Int)] -> Int
sizeOfStaticArray Int
kind [(Int, Int)]
dims'
  Type
ty -> Type -> Maybe Int
getTypeKind Type
ty

-- | Given a static array's 'kind' and 'dimension', calculate its size
sizeOfStaticArray :: Int -> [(Int, Int)] -> Int
sizeOfStaticArray :: Int -> [(Int, Int)] -> Int
sizeOfStaticArray Int
kind' [(Int, Int)]
dimension' =
  let arraySize :: Int
arraySize = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Int
acc (Int
l, Int
h) -> Int
acc forall a. Num a => a -> a -> a
* (Int
h forall a. Num a => a -> a -> a
- Int
l forall a. Num a => a -> a -> a
+ Int
1)) Int
1 [(Int, Int)]
dimension'
  in  Int
kind' 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 = forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey StorageTable -> MemoryBlockName -> SymbolTableEntry -> StorageTable
f forall k a. Map k a
M.empty
 where
  f :: StorageTable -> Name -> SymbolTableEntry -> StorageTable
  f :: StorageTable -> MemoryBlockName -> SymbolTableEntry -> StorageTable
f StorageTable
storageTable MemoryBlockName
symbol SymbolTableEntry
entry = case SymbolTableEntry
entry of
    SVariable Type
ty Location
_ ->
      let mSize :: Maybe Int
mSize = Type -> Maybe Int
getTypeSize' Type
ty
          block :: MemoryBlock
block = MemoryBlock
            { blockSize :: Maybe Int
blockSize    = Maybe Int
mSize
            , storageClass :: StorageClass
storageClass = case Maybe Int
mSize of
                               Maybe Int
Nothing -> StorageClass
Automatic
                               Maybe Int
_       -> StorageClass
Unspecified
            , variables :: [MemoryBlockName]
variables    = [MemoryBlockName
symbol]
            }
      in  forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert MemoryBlockName
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)) <- forall a. Data a => ProgramUnit a -> [Statement a]
allPUS ProgramUnit (Analysis a)
pu ]
      mergeCommonVariables :: Map MemoryBlockName [Declarator (Analysis a)]
-> CommonGroup (Analysis a)
-> Map MemoryBlockName [Declarator (Analysis a)]
mergeCommonVariables Map MemoryBlockName [Declarator (Analysis a)]
mapping (CommonGroup Analysis a
_ SrcSpan
_ Maybe (Expression (Analysis a))
commName AList Declarator (Analysis a)
decls) =
          let commonName :: MemoryBlockName
commonName = case Maybe (Expression (Analysis a))
commName of
                Just Expression (Analysis a)
e  -> MemoryBlockName
"/" forall a. [a] -> [a] -> [a]
++ forall a. Expression (Analysis a) -> MemoryBlockName
srcName Expression (Analysis a)
e forall a. [a] -> [a] -> [a]
++ MemoryBlockName
"/"
                Maybe (Expression (Analysis a))
Nothing -> MemoryBlockName
"*blank_common*"
              precedingDecls :: [Declarator (Analysis a)]
precedingDecls = forall a. a -> Maybe a -> a
fromMaybe [] (forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup MemoryBlockName
commonName Map MemoryBlockName [Declarator (Analysis a)]
mapping)
          in  forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert MemoryBlockName
commonName ([Declarator (Analysis a)]
precedingDecls forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Declarator (Analysis a)
decls) Map MemoryBlockName [Declarator (Analysis a)]
mapping
      commons :: Map MemoryBlockName [Declarator (Analysis a)]
commons = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {a}.
Map MemoryBlockName [Declarator (Analysis a)]
-> CommonGroup (Analysis a)
-> Map MemoryBlockName [Declarator (Analysis a)]
mergeCommonVariables forall k a. Map k a
M.empty (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[CommonGroup (Analysis a)]]
commonGrps)
      processComm :: MemoryBlockName
-> [Declarator (Analysis a)]
-> ProgramUnitModel
-> ProgramUnitModel
processComm MemoryBlockName
commonName [Declarator (Analysis a)]
varDecls (SymbolTable
symTable, StorageTable
mbs) =
          let varExps :: [Expression (Analysis a)]
varExps            = forall a b. (a -> b) -> [a] -> [b]
map forall {a}. Declarator a -> Expression a
declExpr [Declarator (Analysis a)]
varDecls
              varLocations :: [Location]
varLocations       = forall a b. (a -> b) -> [a] -> [b]
map (forall a.
Data a =>
SymbolTable -> Expression (Analysis a) -> Location
getStartLocation SymbolTable
symTable) [Expression (Analysis a)]
varExps
              varSizes :: [Int]
varSizes           = forall a b. (a -> b) -> [a] -> [b]
map (forall a. Data a => SymbolTable -> Expression (Analysis a) -> Int
getSize SymbolTable
symTable) [Expression (Analysis a)]
varExps
              varAccumSizes :: [Int]
varAccumSizes      = forall a. (a -> a -> a) -> [a] -> [a]
scanl1 forall a. Num a => a -> a -> a
(+) [Int]
varSizes
              commBlockLocations :: [Location]
commBlockLocations = forall a b. (a -> b) -> [a] -> [b]
map (MemoryBlockName
commonName, ) (Int
0 forall a. a -> [a] -> [a]
: [Int]
varAccumSizes)
              mbs' :: StorageTable
mbs'               = case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup MemoryBlockName
commonName StorageTable
mbs of
                Just MemoryBlock
_ -> StorageTable
mbs
                Maybe MemoryBlock
Nothing ->
                  let newBlock :: MemoryBlock
newBlock = MemoryBlock { blockSize :: Maybe Int
blockSize = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
varSizes
                                             , storageClass :: StorageClass
storageClass = StorageClass
Common
                                             , variables :: [MemoryBlockName]
variables = []
                                             }
                  in  forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert MemoryBlockName
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  forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ProgramUnitModel -> (Location, Location) -> ProgramUnitModel
f (SymbolTable
symTable, StorageTable
mbs') (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  forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
M.foldrWithKey forall {a}.
Data a =>
MemoryBlockName
-> [Declarator (Analysis a)]
-> ProgramUnitModel
-> ProgramUnitModel
processComm ProgramUnitModel
puModel Map MemoryBlockName [Declarator (Analysis a)]
commons