module Language.Fortran.Vars.MemoryLocation
  ( getLocation
  , generateLinearizedIndexRange
  , getStartLocation
  )
where

import           Data.Data                      ( Data )
import           Data.List                      ( foldl' )
import qualified Data.Map                      as M
import           Language.Fortran.Analysis      ( Analysis
                                                , srcName
                                                )
import           Language.Fortran.AST           ( AList(..)
                                                , Expression(..)
                                                , Index(..)
                                                , Name
                                                , Value(..)
                                                )

import           Language.Fortran.Vars.Eval
                                                ( eval )
import           Language.Fortran.Vars.Kind
                                                ( toInt
                                                , getTypeKind
                                                )
import           Language.Fortran.Vars.Range
                                                ( Range )
import           Language.Fortran.Vars.Types
                                                ( SymbolTableEntry(..)
                                                , Type(..)
                                                , SemType(..)
                                                , Location
                                                , Offset
                                                , SymbolTable
                                                )


isIxSingle :: Index a -> Bool
isIxSingle :: forall a. Index a -> Bool
isIxSingle IxSingle{} = Bool
True
isIxSingle IxRange{}  = Bool
False

linearizedIndex :: [Int] -> [(Int, Int)] -> Int
linearizedIndex :: [Offset] -> [(Offset, Offset)] -> Offset
linearizedIndex [Offset]
indices [(Offset, Offset)]
dimensions =
  -- The normalized index starts at 0
  let normalizedIndices :: [Offset]
normalizedIndices =
          (Offset -> (Offset, Offset) -> Offset)
-> [Offset] -> [(Offset, Offset)] -> [Offset]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (((Offset, (Offset, Offset)) -> Offset)
-> Offset -> (Offset, Offset) -> Offset
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (\(Offset
i, (Offset
lower, Offset
_)) -> Offset
i Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
- Offset
lower)) [Offset]
indices [(Offset, Offset)]
dimensions
      dimSizes :: [Offset]
dimSizes    = ((Offset, Offset) -> Offset) -> [(Offset, Offset)] -> [Offset]
forall a b. (a -> b) -> [a] -> [b]
map (\(Offset
lower, Offset
upper) -> Offset
upper Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
- Offset
lower Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
+ Offset
1) [(Offset, Offset)]
dimensions
      leadingDims :: [Offset]
leadingDims = [Offset] -> [Offset]
calcLeadingDimensions [Offset]
dimSizes
  in  (Offset -> (Offset, Offset) -> Offset)
-> Offset -> [(Offset, Offset)] -> Offset
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Offset
acc (Offset
i, Offset
ld) -> Offset
acc Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
+ Offset
i Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
* Offset
ld)
             Offset
0
             ([Offset] -> [Offset] -> [(Offset, Offset)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Offset]
normalizedIndices [Offset]
leadingDims)
 where
  calcLeadingDimensions :: [Int] -> [Int]
  calcLeadingDimensions :: [Offset] -> [Offset]
calcLeadingDimensions []       = []
  calcLeadingDimensions [Offset]
dimSizes = [Offset] -> [Offset] -> [Offset]
forall {a}. Num a => [a] -> [a] -> [a]
calc [Offset
1] [Offset]
dimSizes
   where
    calc :: [a] -> [a] -> [a]
calc [a]
_   []      = []
    calc [a]
acc [a
_    ] = [a] -> [a]
forall a. [a] -> [a]
reverse [a]
acc
    calc [a]
acc (a
h : [a]
t) = [a] -> [a] -> [a]
calc (([a] -> a
forall a. [a] -> a
head [a]
acc a -> a -> a
forall a. Num a => a -> a -> a
* a
h) a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc) [a]
t

-- | Given only single indices return the 'Range' in memory that
-- these indices point to.
generateLinearizedIndexRange :: [Int] -> Int -> [(Int, Int)] -> Int -> Range
generateLinearizedIndexRange :: [Offset]
-> Offset -> [(Offset, Offset)] -> Offset -> (Offset, Offset)
generateLinearizedIndexRange [Offset]
intIndices Offset
start [(Offset, Offset)]
dims Offset
kind =
  let offset :: Offset
offset = [Offset] -> [(Offset, Offset)] -> Offset
linearizedIndex [Offset]
intIndices [(Offset, Offset)]
dims Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
* Offset
kind
  in  (Offset
start Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
+ Offset
offset, Offset
start Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
+ Offset
offset Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
+ Offset
kind Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
- Offset
1)

findBlockOffset :: SymbolTable -> Name -> Offset -> Location
findBlockOffset :: SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable Name
symbol Offset
offset = case Name -> SymbolTable -> Maybe SymbolTableEntry
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
symbol SymbolTable
symTable of
  Just (SVariable Type
_ (Name
blockName, Offset
start)) -> (Name
blockName, Offset
start Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
+ Offset
offset)
  Just SymbolTableEntry
entry -> Name -> Location
forall a. HasCallStack => Name -> a
error (Name -> Location) -> Name -> Location
forall a b. (a -> b) -> a -> b
$ Name
symbol Name -> Name -> Name
forall a. [a] -> [a] -> [a]
++ Name
" is not a variable - " Name -> Name -> Name
forall a. [a] -> [a] -> [a]
++ SymbolTableEntry -> Name
forall a. Show a => a -> Name
show SymbolTableEntry
entry
  Maybe SymbolTableEntry
Nothing ->
    Name -> Location
forall a. HasCallStack => Name -> a
error
      (Name -> Location) -> Name -> Location
forall a b. (a -> b) -> a -> b
$  Name
"Unable to find location for symbol "
      Name -> Name -> Name
forall a. [a] -> [a] -> [a]
++ Name
symbol
      Name -> Name -> Name
forall a. [a] -> [a] -> [a]
++ Name
" at offset "
      Name -> Name -> Name
forall a. [a] -> [a] -> [a]
++ Offset -> Name
forall a. Show a => a -> Name
show Offset
offset

calculateOffset :: Data a => SymbolTable -> Name -> [Index (Analysis a)] -> Int
-- array index c(2,4)
calculateOffset :: forall a.
Data a =>
SymbolTable -> Name -> [Index (Analysis a)] -> Offset
calculateOffset SymbolTable
symTable Name
symbol indices :: [Index (Analysis a)]
indices@(IxSingle{} : [Index (Analysis a)]
_) =
  let Just SymbolTableEntry
entry = Name -> SymbolTable -> Maybe SymbolTableEntry
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
symbol SymbolTable
symTable
  in  case SymbolTableEntry
entry of
        SVariable (TArray Type
ty (Just [(Offset, Offset)]
dims)) Location
_ ->
          let ixSingles :: [Index (Analysis a)]
ixSingles    = (Index (Analysis a) -> Bool)
-> [Index (Analysis a)] -> [Index (Analysis a)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Index (Analysis a) -> Bool
forall a. Index a -> Bool
isIxSingle [Index (Analysis a)]
indices
              Just Offset
kind    = Type -> Maybe Offset
getTypeKind Type
ty
              arrayIndices :: [Offset]
arrayIndices = (Index (Analysis a) -> Offset) -> [Index (Analysis a)] -> [Offset]
forall a b. (a -> b) -> [a] -> [b]
map Index (Analysis a) -> Offset
forall {a}. Index a -> Offset
toIndices [Index (Analysis a)]
ixSingles
                 where
                  toIndices :: Index a -> Offset
toIndices (IxSingle a
_ SrcSpan
_ Maybe Name
_ Expression a
expr) = ExpVal -> Offset
toInt (ExpVal -> Offset) -> ExpVal -> Offset
forall a b. (a -> b) -> a -> b
$ SymbolTable -> Expression a -> ExpVal
forall a. SymbolTable -> Expression a -> ExpVal
eval SymbolTable
symTable Expression a
expr
                  toIndices Index a
_ = Name -> Offset
forall a. HasCallStack => Name -> a
error Name
"toIndices: unexpected input"
          in  [Offset] -> [(Offset, Offset)] -> Offset
linearizedIndex [Offset]
arrayIndices [(Offset, Offset)]
dims Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
* Offset
kind
        SymbolTableEntry
_ ->
          Name -> Offset
forall a. HasCallStack => Name -> a
error Name
"Only array-typed VariableEntries are expected at this point"
-- substring c(:5)
calculateOffset SymbolTable
_ Name
_ (IxRange Analysis a
_ SrcSpan
_ Maybe (Expression (Analysis a))
Nothing Maybe (Expression (Analysis a))
_ Maybe (Expression (Analysis a))
_ : [Index (Analysis a)]
_) = Offset
0
-- substring c(5:)
calculateOffset SymbolTable
symTable Name
_ (IxRange Analysis a
_ SrcSpan
_ (Just Expression (Analysis a)
lowerIndex) Maybe (Expression (Analysis a))
_ Maybe (Expression (Analysis a))
_ : [Index (Analysis a)]
_) =
  ExpVal -> Offset
toInt (SymbolTable -> Expression (Analysis a) -> ExpVal
forall a. SymbolTable -> Expression a -> ExpVal
eval SymbolTable
symTable Expression (Analysis a)
lowerIndex) Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
- Offset
1
calculateOffset SymbolTable
_ Name
_ [Index (Analysis a)]
_ = Name -> Offset
forall a. HasCallStack => Name -> a
error Name
"calculateOffset: invalid index"

-- | Given a 'SymbolTable' and some 'Expression' (which is assumed to have been predetermined
-- to be of some variable type), return the 'Location' that the variable in question will be
-- located in memory
getLocation :: Data a => SymbolTable -> Expression (Analysis a) -> Location
-- variable
getLocation :: forall a.
Data a =>
SymbolTable -> Expression (Analysis a) -> Location
getLocation SymbolTable
symTable e :: Expression (Analysis a)
e@(ExpValue Analysis a
_ SrcSpan
_ (ValVariable Name
_)) =
  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable (Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e) Offset
0
-- array index c(2,4)
-- substring c(5:10)
getLocation SymbolTable
symTable (ExpSubscript Analysis a
_ SrcSpan
_ e :: Expression (Analysis a)
e@ExpValue{} (AList Analysis a
_ SrcSpan
_ [Index (Analysis a)]
indices)) =
  let symbol :: Name
symbol = Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e
      offset :: Offset
offset = SymbolTable -> Name -> [Index (Analysis a)] -> Offset
forall a.
Data a =>
SymbolTable -> Name -> [Index (Analysis a)] -> Offset
calculateOffset SymbolTable
symTable Name
symbol [Index (Analysis a)]
indices
  in  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable Name
symbol Offset
offset
-- array index and substring c(2,4)(1:20)
getLocation SymbolTable
symTable (ExpSubscript Analysis a
_ SrcSpan
_ (ExpSubscript Analysis a
_ SrcSpan
_ e :: Expression (Analysis a)
e@ExpValue{} (AList Analysis a
_ SrcSpan
_ [Index (Analysis a)]
indices)) (AList Analysis a
_ SrcSpan
_ [Index (Analysis a)]
subs))
  = let symbol :: Name
symbol = Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e
        offset :: Offset
offset =
            SymbolTable -> Name -> [Index (Analysis a)] -> Offset
forall a.
Data a =>
SymbolTable -> Name -> [Index (Analysis a)] -> Offset
calculateOffset SymbolTable
symTable Name
symbol [Index (Analysis a)]
indices
              Offset -> Offset -> Offset
forall a. Num a => a -> a -> a
+ SymbolTable -> Name -> [Index (Analysis a)] -> Offset
forall a.
Data a =>
SymbolTable -> Name -> [Index (Analysis a)] -> Offset
calculateOffset SymbolTable
symTable Name
symbol [Index (Analysis a)]
subs
    in  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable Name
symbol Offset
offset
-- array within common block with dimensions declaration: common /block/ a, b(10)
getLocation SymbolTable
symTable (ExpFunctionCall Analysis a
_ SrcSpan
_ e :: Expression (Analysis a)
e@ExpValue{} Maybe (AList Argument (Analysis a))
_) =
  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable (Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e) Offset
0
getLocation SymbolTable
_ Expression (Analysis a)
_ = Name -> Location
forall a. HasCallStack => Name -> a
error Name
"getLocation : Not a variable expression"

-- | Given a 'SymbolTable' and some 'Expression' (which is assumed to have been
-- predetermined to be of some variable type), return the start 'Location' that
-- the variable in question will be located in memory.
-- Start 'Location' is the begining of greater data structure that a variable
-- belongs to e.g. start 'Location' of c(20) is 'Location' of c
getStartLocation :: Data a => SymbolTable -> Expression (Analysis a) -> Location
-- variable
getStartLocation :: forall a.
Data a =>
SymbolTable -> Expression (Analysis a) -> Location
getStartLocation SymbolTable
symTable e :: Expression (Analysis a)
e@(ExpValue Analysis a
_ SrcSpan
_ (ValVariable Name
_)) =
  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable (Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e) Offset
0
-- dimensions specification within COMMON c(10,20)
getStartLocation SymbolTable
symTable (ExpSubscript Analysis a
_ SrcSpan
_ e :: Expression (Analysis a)
e@ExpValue{} AList Index (Analysis a)
_) =
  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable (Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e) Offset
0
-- dimensions specification within COMMON c(10,20)
--   (vars with nonstandard kind declared after COMMON block)
getStartLocation SymbolTable
symTable (ExpFunctionCall Analysis a
_ SrcSpan
_ e :: Expression (Analysis a)
e@ExpValue{} Maybe (AList Argument (Analysis a))
_) =
  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable (Expression (Analysis a) -> Name
forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e) Offset
0
getStartLocation SymbolTable
_ Expression (Analysis a)
_ = Name -> Location
forall a. HasCallStack => Name -> a
error Name
"getStartLocation : Not a variable expression"