{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Language.Fortran.Vars.Assignments
  ( allAssignStmts
  )
where

import           Data.Data                      ( Data )
import           Data.List                      ( foldl' )
import qualified Data.Map                      as M
import           Data.Generics.Uniplate.Data    ( universeBi )

import           Language.Fortran.Analysis      ( Analysis )
import           Language.Fortran.AST           ( ProgramUnit
                                                , Statement(..)
                                                , DataGroup(..)
                                                , Expression(..)
                                                , Declarator(..)
                                                , DeclaratorType(..)
                                                , Value(..)
                                                , aStrip
                                                )
import           Language.Fortran.Util.Position ( Spanned(..) )

import           Language.Fortran.Extras.Encoding
                                                ( pprint77l )
import           Language.Fortran.Vars.StructureTable
                                                ( collectStructures )
import           Language.Fortran.Vars.SymbolTable
                                                ( collectSymbols )
import           Language.Fortran.Vars.Types    ( SymbolTable
                                                , StructureTable
                                                , SymbolTableEntry(..)
                                                , Dims(..)
                                                , dimsTraverse
                                                , Dim(..)
                                                , Type
                                                , SemType(..)
                                                , TypeError(..)
                                                , typeError
                                                )
import           Language.Fortran.Vars.TypeCheck
                                                ( typeOf )

-- | Method to retrieve the type of the lhs and expression on the rhs of all
-- assign like statements: expression assign, parameter, data and declarations
-- It returns the type of the lhs due to the expansion of array types to scalar
-- types in datagroups and declarations and returns TypeError's for the
-- expressions it can't calculate.
allAssignStmts
  :: forall a
   . Data a
  => ProgramUnit (Analysis a)
  -> [Either TypeError (Type, Expression (Analysis a))]
allAssignStmts :: forall a.
Data a =>
ProgramUnit (Analysis a)
-> [Either TypeError (Type, Expression (Analysis a))]
allAssignStmts ProgramUnit (Analysis a)
pu =
  let
    symt :: SymbolTable
symt = forall a. Data a => ProgramUnit (Analysis a) -> SymbolTable
collectSymbols ProgramUnit (Analysis a)
pu
    strt :: StructureTable
strt = forall a.
Data a =>
SymbolTable -> ProgramUnit (Analysis a) -> StructureTable
collectStructures SymbolTable
symt ProgramUnit (Analysis a)
pu
  in
    [ (, Expression (Analysis a)
e) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either TypeError Type
ty
    | StExpressionAssign Analysis a
_ SrcSpan
_ Expression (Analysis a)
v Expression (Analysis a)
e <- forall from to. Biplate from to => from -> [to]
universeBi ProgramUnit (Analysis a)
pu :: [Statement (Analysis a)]
    , let ty :: Either TypeError Type
ty = forall a.
StructureTable
-> SymbolTable -> Expression a -> Either TypeError Type
typeOf StructureTable
strt SymbolTable
symt Expression (Analysis a)
v
    ]
    forall a. Semigroup a => a -> a -> a
<> [ (, Expression (Analysis a)
e) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either TypeError Type
ty
       | StParameter Analysis a
_ SrcSpan
_ AList Declarator (Analysis a)
decls <- forall from to. Biplate from to => from -> [to]
universeBi ProgramUnit (Analysis a)
pu :: [Statement (Analysis a)]
       , Declarator Analysis a
_ SrcSpan
_ Expression (Analysis a)
v DeclaratorType (Analysis a)
ScalarDecl Maybe (Expression (Analysis a))
_ (Just Expression (Analysis a)
e) <- forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Declarator (Analysis a)
decls
       , let ty :: Either TypeError Type
ty = forall a.
StructureTable
-> SymbolTable -> Expression a -> Either TypeError Type
typeOf StructureTable
strt SymbolTable
symt Expression (Analysis a)
v
       ]
    forall a. Semigroup a => a -> a -> a
<> [ Either TypeError (Type, Expression (Analysis a))
res
       | StData Analysis a
_ SrcSpan
_ AList DataGroup (Analysis a)
groups <- forall from to. Biplate from to => from -> [to]
universeBi ProgramUnit (Analysis a)
pu :: [Statement (Analysis a)]
       , Either TypeError (Type, Expression (Analysis a))
res               <- forall a.
StructureTable
-> SymbolTable
-> [DataGroup (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
dataGroups StructureTable
strt SymbolTable
symt (forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList DataGroup (Analysis a)
groups)
       ]
    forall a. Semigroup a => a -> a -> a
<> [ Either TypeError (Type, Expression (Analysis a))
res
       | StDeclaration Analysis a
_ SrcSpan
_ TypeSpec (Analysis a)
_ Maybe (AList Attribute (Analysis a))
_ AList Declarator (Analysis a)
decls <-
         forall from to. Biplate from to => from -> [to]
universeBi ProgramUnit (Analysis a)
pu :: [Statement (Analysis a)]
       , Either TypeError (Type, Expression (Analysis a))
res <- forall a.
StructureTable
-> SymbolTable
-> [Declarator (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
declarators StructureTable
strt SymbolTable
symt forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Declarator (Analysis a)
decls
       ]

-- | Deal with data groups
dataGroups
  :: StructureTable
  -> SymbolTable
  -> [DataGroup (Analysis a)]
  -> [Either TypeError (Type, Expression (Analysis a))]
dataGroups :: forall a.
StructureTable
-> SymbolTable
-> [DataGroup (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
dataGroups StructureTable
strt SymbolTable
symt = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DataGroup (Analysis a)
-> [Either TypeError (Type, Expression (Analysis a))]
f
 where
  f :: DataGroup (Analysis a)
-> [Either TypeError (Type, Expression (Analysis a))]
f (DataGroup Analysis a
_ SrcSpan
_ AList Expression (Analysis a)
vas AList Expression (Analysis a)
eas) =
    let vs :: [Expression (Analysis a)]
vs         = forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Expression (Analysis a)
vas
        es :: [Expression (Analysis a)]
es         = forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Expression (Analysis a)
eas
        expandedVs :: [Either TypeError Type]
expandedVs = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a.
StructureTable
-> SymbolTable
-> Expression (Analysis a)
-> [Either TypeError Type]
expandArrays StructureTable
strt SymbolTable
symt) [Expression (Analysis a)]
vs
        g :: Either a a -> b -> Either a (a, b)
g (Left  a
err) b
_ = forall a b. a -> Either a b
Left a
err
        g (Right a
ty ) b
e = forall a b. b -> Either a b
Right (a
ty, b
e)
    in  forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {a} {a} {b}. Either a a -> b -> Either a (a, b)
g [Either TypeError Type]
expandedVs [Expression (Analysis a)]
es

-- | Expands declarators to lhs type and rhs expression
declarators
  :: StructureTable
  -> SymbolTable
  -> [Declarator (Analysis a)]
  -> [Either TypeError (Type, Expression (Analysis a))]
declarators :: forall a.
StructureTable
-> SymbolTable
-> [Declarator (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
declarators StructureTable
strt SymbolTable
symt = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Declarator (Analysis a)
-> [Either TypeError (Type, Expression (Analysis a))]
f where
  f :: Declarator (Analysis a)
-> [Either TypeError (Type, Expression (Analysis a))]
f (Declarator Analysis a
_ SrcSpan
_ Expression (Analysis a)
v DeclaratorType (Analysis a)
ScalarDecl Maybe (Expression (Analysis a))
_ (Just Expression (Analysis a)
e)) =
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (, Expression (Analysis a)
e) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a.
StructureTable
-> SymbolTable -> Expression a -> Either TypeError Type
typeOf StructureTable
strt SymbolTable
symt Expression (Analysis a)
v
  f d :: Declarator (Analysis a)
d@(Declarator Analysis a
_ SrcSpan
_ (ExpValue Analysis a
_ SrcSpan
s (ValVariable Name
v)) ArrayDecl{} Maybe (Expression (Analysis a))
_ (Just (ExpInitialisation Analysis a
_ SrcSpan
_ AList Expression (Analysis a)
vals)))
    = case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
v SymbolTable
symt of
      Just (SVariable (TArray Type
ty Dimensions
dims') Location
_) ->
       case forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
Dims t (f a) -> f (Dims t a)
dimsTraverse Dimensions
dims' of -- only static arrays
       Just (DimsExplicitShape NonEmpty (Dim Int)
dims) ->
        let tys :: [Type]
tys   = forall (t :: * -> *) a. Foldable t => t (Dim Int) -> a -> [a]
expandDimensions NonEmpty (Dim Int)
dims Type
ty
            vals' :: [Expression (Analysis a)]
vals' = forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Expression (Analysis a)
vals
        in  if forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
tys forall a. Eq a => a -> a -> Bool
/= forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expression (Analysis a)]
vals'
              then
                forall (f :: * -> *) a. Applicative f => a -> f a
pure
                forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> Either a b
Left
                forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcSpan -> Name -> TypeError
typeError SrcSpan
s
                forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Monoid a => [a] -> a
mconcat
                forall a b. (a -> b) -> a -> b
$ [ Name
"Length of lhs and rhs in declarator do not match: "
                  , Name
": "
                  , forall a. IndentablePretty a => a -> Name
pprint77l Declarator (Analysis a)
d
                  ]
              else forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall (t :: * -> *) a. Foldable t => t (Dim Int) -> a -> [a]
expandDimensions NonEmpty (Dim Int)
dims Type
ty) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Expression (Analysis a)
vals
       -- only static explicit-shape arrays permitted
       Just{} ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure
          forall b c a. (b -> c) -> (a -> b) -> a -> c
.  forall a b. a -> Either a b
Left
          forall b c a. (b -> c) -> (a -> b) -> a -> c
.  SrcSpan -> Name -> TypeError
typeError SrcSpan
s
          forall a b. (a -> b) -> a -> b
$  Name
"Unexpected lhs in array declaration at: "
          forall a. Semigroup a => a -> a -> a
<> forall a. IndentablePretty a => a -> Name
pprint77l Declarator (Analysis a)
d
       Maybe (Dims NonEmpty Int)
Nothing ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure
          forall b c a. (b -> c) -> (a -> b) -> a -> c
.  forall a b. a -> Either a b
Left
          forall b c a. (b -> c) -> (a -> b) -> a -> c
.  SrcSpan -> Name -> TypeError
typeError SrcSpan
s
          forall a b. (a -> b) -> a -> b
$  Name
"Unexpected lhs in array declaration at: "
          forall a. Semigroup a => a -> a -> a
<> forall a. IndentablePretty a => a -> Name
pprint77l Declarator (Analysis a)
d
      Maybe SymbolTableEntry
_ ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure
          forall b c a. (b -> c) -> (a -> b) -> a -> c
.  forall a b. a -> Either a b
Left
          forall b c a. (b -> c) -> (a -> b) -> a -> c
.  SrcSpan -> Name -> TypeError
typeError SrcSpan
s
          forall a b. (a -> b) -> a -> b
$  Name
"Unexpected lhs in array declaration at: "
          forall a. Semigroup a => a -> a -> a
<> forall a. IndentablePretty a => a -> Name
pprint77l Declarator (Analysis a)
d
  f Declarator (Analysis a)
_ = []  -- All other cases have no initial values

-- | Expands arrays to give a list of types given the length of an array
expandArrays
  :: StructureTable
  -> SymbolTable
  -> Expression (Analysis a)
  -> [Either TypeError Type]
expandArrays :: forall a.
StructureTable
-> SymbolTable
-> Expression (Analysis a)
-> [Either TypeError Type]
expandArrays StructureTable
strt SymbolTable
symt Expression (Analysis a)
e = case Expression (Analysis a)
e of
  ExpValue Analysis a
_ SrcSpan
_ (ValVariable Name
var) -> case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
var SymbolTable
symt of
    Just (SVariable (TArray Type
sty Dimensions
dims') Location
_) ->
     case forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
Dims t (f a) -> f (Dims t a)
dimsTraverse Dimensions
dims' of
      Just (DimsExplicitShape NonEmpty (Dim Int)
dims) -> forall (t :: * -> *) a. Foldable t => t (Dim Int) -> a -> [a]
expandDimensions NonEmpty (Dim Int)
dims (forall a b. b -> Either a b
Right Type
sty)
      Maybe (Dims NonEmpty Int)
_ -> [forall a b. b -> Either a b
Right Type
sty]
    Just (SVariable Type
ty Location
_) -> [forall a b. b -> Either a b
Right Type
ty]
    Maybe SymbolTableEntry
_ ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure
        forall b c a. (b -> c) -> (a -> b) -> a -> c
.  forall a b. a -> Either a b
Left
        forall b c a. (b -> c) -> (a -> b) -> a -> c
.  SrcSpan -> Name -> TypeError
typeError (forall a. Spanned a => a -> SrcSpan
getSpan Expression (Analysis a)
e)
        forall a b. (a -> b) -> a -> b
$  Name
"Got unexpected lhs type: "
        forall a. Semigroup a => a -> a -> a
<> forall a. IndentablePretty a => a -> Name
pprint77l Expression (Analysis a)
e
  Expression (Analysis a)
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a.
StructureTable
-> SymbolTable -> Expression a -> Either TypeError Type
typeOf StructureTable
strt SymbolTable
symt Expression (Analysis a)
e

-- | Function to expand dimensions into appropriate number of types for use in
-- other expand functions
expandDimensions :: Foldable t => t (Dim Int) -> a -> [a]
expandDimensions :: forall (t :: * -> *) a. Foldable t => t (Dim Int) -> a -> [a]
expandDimensions t (Dim Int)
dims =
  forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
acc (Dim Int
lb Int
ub) -> forall a. Num a => a -> a
abs (Int
ub forall a. Num a => a -> a -> a
- Int
lb forall a. Num a => a -> a -> a
+ Int
1) forall a. Num a => a -> a -> a
* Int
acc) Int
1 t (Dim Int)
dims)