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

import qualified Data.Foldable                 as Foldable
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(..)
                                                , SemType(..)
                                                , Location
                                                , Offset
                                                , SymbolTable
                                                , Dim(..), Dimensions
                                                , dimensionsToTuples
                                                )

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 =
          forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall a b c. ((a, b) -> c) -> a -> b -> c
curry (\(Offset
i, (Offset
lower, Offset
_)) -> Offset
i forall a. Num a => a -> a -> a
- Offset
lower)) [Offset]
indices [(Offset, Offset)]
dimensions
      dimSizes :: [Offset]
dimSizes    = forall a b. (a -> b) -> [a] -> [b]
map (\(Offset
lower, Offset
upper) -> Offset
upper forall a. Num a => a -> a -> a
- Offset
lower forall a. Num a => a -> a -> a
+ Offset
1) [(Offset, Offset)]
dimensions
      leadingDims :: [Offset]
leadingDims = [Offset] -> [Offset]
calcLeadingDimensions [Offset]
dimSizes
  in  forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Offset
acc (Offset
i, Offset
ld) -> Offset
acc forall a. Num a => a -> a -> a
+ Offset
i forall a. Num a => a -> a -> a
* Offset
ld)
             Offset
0
             (forall a b. [a] -> [b] -> [(a, b)]
zip [Offset]
normalizedIndices [Offset]
leadingDims)
 where
  calcLeadingDimensions :: [Int] -> [Int]
  calcLeadingDimensions :: [Offset] -> [Offset]
calcLeadingDimensions []       = []
  calcLeadingDimensions [Offset]
dimSizes = forall {a}. Num a => [a] -> [a] -> [a]
calc [Offset
1] [Offset]
dimSizes
   where
    calc :: [a] -> [a] -> [a]
calc [a]
_   []      = []
    calc [a]
acc [a
_    ] = forall a. [a] -> [a]
reverse [a]
acc
    calc [a]
acc (a
h : [a]
t) = [a] -> [a] -> [a]
calc ((forall a. [a] -> a
head [a]
acc forall a. Num a => a -> a -> a
* a
h) forall a. a -> [a] -> [a]
: [a]
acc) [a]
t

-- | Given only single indices return the 'Range' in memory that
-- these indices point to.
generateLinearizedIndexRange
    :: (Functor t, Foldable t) => [Int] -> Int -> t (Dim Int) -> Int -> Range
generateLinearizedIndexRange :: forall (t :: * -> *).
(Functor t, Foldable t) =>
[Offset] -> Offset -> t (Dim Offset) -> Offset -> (Offset, Offset)
generateLinearizedIndexRange [Offset]
intIndices Offset
start t (Dim Offset)
dims Offset
kind =
  let offset :: Offset
offset = [Offset] -> [(Offset, Offset)] -> Offset
linearizedIndex [Offset]
intIndices [(Offset, Offset)]
dims' forall a. Num a => a -> a -> a
* Offset
kind
  in  (Offset
start forall a. Num a => a -> a -> a
+ Offset
offset, Offset
start forall a. Num a => a -> a -> a
+ Offset
offset forall a. Num a => a -> a -> a
+ Offset
kind forall a. Num a => a -> a -> a
- Offset
1)
  where
    -- TODO Ideally, we stay in our foldable for as long as possible. Shift this
    -- into 'linearizedIndex' and try.
    dims' :: [(Offset, Offset)]
dims' = forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Dim Offset
lb Offset
ub) -> (Offset
lb, Offset
ub)) t (Dim Offset)
dims

findBlockOffset :: SymbolTable -> Name -> Offset -> Location
findBlockOffset :: SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable Name
symbol Offset
offset = case 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 forall a. Num a => a -> a -> a
+ Offset
offset)
  Just SymbolTableEntry
entry -> forall a. HasCallStack => Name -> a
error forall a b. (a -> b) -> a -> b
$ Name
symbol forall a. [a] -> [a] -> [a]
++ Name
" is not a variable - " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> Name
show SymbolTableEntry
entry
  Maybe SymbolTableEntry
Nothing ->
    forall a. HasCallStack => Name -> a
error
      forall a b. (a -> b) -> a -> b
$  Name
"Unable to find location for symbol "
      forall a. [a] -> [a] -> [a]
++ Name
symbol
      forall a. [a] -> [a] -> [a]
++ Name
" at offset "
      forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> Name
show Offset
offset

calculateOffset
  :: Data a => SymbolTable -> Name -> [Index (Analysis a)] -> Maybe Int
-- array index c(2,4)
calculateOffset :: forall a.
Data a =>
SymbolTable -> Name -> [Index (Analysis a)] -> Maybe Offset
calculateOffset SymbolTable
symTable Name
symbol indices :: [Index (Analysis a)]
indices@(IxSingle{} : [Index (Analysis a)]
_) =
  let Just SymbolTableEntry
entry = 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 Dimensions
dims) Location
_ ->
        case Dimensions -> Maybe [(Offset, Offset)]
dimensionsToTuples Dimensions
dims of
          Maybe [(Offset, Offset)]
Nothing -> forall a. HasCallStack => Name -> a
error Name
"expected a static array, got dynamic"
          Just [(Offset, Offset)]
dims' ->
            let
              ixSingles :: [Index (Analysis a)]
ixSingles    = forall a. (a -> Bool) -> [a] -> [a]
takeWhile forall a. Index a -> Bool
isIxSingle [Index (Analysis a)]
indices
              Just Offset
kind    = Type -> Maybe Offset
getTypeKind Type
ty
              arrayIndices :: Maybe [Offset]
arrayIndices = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const forall a. Maybe a
Nothing) forall a. a -> Maybe a
Just
                forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Index (Analysis a) -> Either Name Offset
toIndices [Index (Analysis a)]
ixSingles
               where
                toIndices :: Index (Analysis a) -> Either Name Offset
toIndices (IxSingle Analysis a
_ SrcSpan
_ Maybe Name
_ Expression (Analysis a)
expr) = ExpVal -> Offset
toInt forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. SymbolTable -> Expression a -> Either Name ExpVal
eval' SymbolTable
symTable Expression (Analysis a)
expr
                toIndices Index (Analysis a)
_ = forall a. HasCallStack => Name -> a
error Name
"toIndices: unexpected input"
            in
              (\[Offset]
x -> [Offset] -> [(Offset, Offset)] -> Offset
linearizedIndex [Offset]
x [(Offset, Offset)]
dims' forall a. Num a => a -> a -> a
* Offset
kind) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe [Offset]
arrayIndices
      SymbolTableEntry
_ -> 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)]
_) = forall a. a -> Maybe a
Just 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)]
_) =
  let val :: Either Name ExpVal
val = forall a. SymbolTable -> Expression a -> Either Name ExpVal
eval' SymbolTable
symTable Expression (Analysis a)
lowerIndex
  in  forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const forall a. Maybe a
Nothing) (\ExpVal
x -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ ExpVal -> Offset
toInt ExpVal
x forall a. Num a => a -> a -> a
- Offset
1) Either Name ExpVal
val
calculateOffset SymbolTable
_ Name
_ [Index (Analysis a)]
_ = 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) -> Maybe Location
-- variable
getLocation :: forall a.
Data a =>
SymbolTable -> Expression (Analysis a) -> Maybe Location
getLocation SymbolTable
symTable e :: Expression (Analysis a)
e@(ExpValue Analysis a
_ SrcSpan
_ (ValVariable Name
_)) =
  forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable (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 = forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e
      offset :: Maybe Offset
offset = forall a.
Data a =>
SymbolTable -> Name -> [Index (Analysis a)] -> Maybe Offset
calculateOffset SymbolTable
symTable Name
symbol [Index (Analysis a)]
indices
  in  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable Name
symbol forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe 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 = forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e
        offset :: Maybe Offset
offset =
            forall a. Num a => a -> a -> a
(+)
              forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
Data a =>
SymbolTable -> Name -> [Index (Analysis a)] -> Maybe Offset
calculateOffset SymbolTable
symTable Name
symbol [Index (Analysis a)]
indices
              forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a.
Data a =>
SymbolTable -> Name -> [Index (Analysis a)] -> Maybe Offset
calculateOffset SymbolTable
symTable Name
symbol [Index (Analysis a)]
subs
    in  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable Name
symbol forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe 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{} AList Argument (Analysis a)
_) =
  forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable (forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e) Offset
0
getLocation SymbolTable
_ Expression (Analysis a)
_ = forall a. Maybe a
Nothing

-- | 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 (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 (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{} AList Argument (Analysis a)
_) =
  SymbolTable -> Name -> Offset -> Location
findBlockOffset SymbolTable
symTable (forall a. Expression (Analysis a) -> Name
srcName Expression (Analysis a)
e) Offset
0
getStartLocation SymbolTable
_ Expression (Analysis a)
_ = forall a. HasCallStack => Name -> a
error Name
"getStartLocation : Not a variable expression"