{-# 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(..)
                                                , Dimensions
                                                , 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 = ProgramUnit (Analysis a) -> SymbolTable
forall a. Data a => ProgramUnit (Analysis a) -> SymbolTable
collectSymbols ProgramUnit (Analysis a)
pu
    strt :: StructureTable
strt = SymbolTable -> ProgramUnit (Analysis a) -> StructureTable
forall a.
Data a =>
SymbolTable -> ProgramUnit (Analysis a) -> StructureTable
collectStructures SymbolTable
symt ProgramUnit (Analysis a)
pu
  in
    [ (, Expression (Analysis a)
e) (Type -> (Type, Expression (Analysis a)))
-> Either TypeError Type
-> Either TypeError (Type, Expression (Analysis a))
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 <- ProgramUnit (Analysis a) -> [Statement (Analysis a)]
forall from to. Biplate from to => from -> [to]
universeBi ProgramUnit (Analysis a)
pu :: [Statement (Analysis a)]
    , let ty :: Either TypeError Type
ty = StructureTable
-> SymbolTable -> Expression (Analysis a) -> Either TypeError Type
forall a.
StructureTable
-> SymbolTable -> Expression a -> Either TypeError Type
typeOf StructureTable
strt SymbolTable
symt Expression (Analysis a)
v
    ]
    [Either TypeError (Type, Expression (Analysis a))]
-> [Either TypeError (Type, Expression (Analysis a))]
-> [Either TypeError (Type, Expression (Analysis a))]
forall a. Semigroup a => a -> a -> a
<> [ (, Expression (Analysis a)
e) (Type -> (Type, Expression (Analysis a)))
-> Either TypeError Type
-> Either TypeError (Type, Expression (Analysis a))
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 <- ProgramUnit (Analysis a) -> [Statement (Analysis a)]
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) <- AList Declarator (Analysis a) -> [Declarator (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Declarator (Analysis a)
decls
       , let ty :: Either TypeError Type
ty = StructureTable
-> SymbolTable -> Expression (Analysis a) -> Either TypeError Type
forall a.
StructureTable
-> SymbolTable -> Expression a -> Either TypeError Type
typeOf StructureTable
strt SymbolTable
symt Expression (Analysis a)
v
       ]
    [Either TypeError (Type, Expression (Analysis a))]
-> [Either TypeError (Type, Expression (Analysis a))]
-> [Either TypeError (Type, Expression (Analysis a))]
forall a. Semigroup a => a -> a -> a
<> [ Either TypeError (Type, Expression (Analysis a))
res
       | StData Analysis a
_ SrcSpan
_ AList DataGroup (Analysis a)
groups <- ProgramUnit (Analysis a) -> [Statement (Analysis a)]
forall from to. Biplate from to => from -> [to]
universeBi ProgramUnit (Analysis a)
pu :: [Statement (Analysis a)]
       , Either TypeError (Type, Expression (Analysis a))
res               <- StructureTable
-> SymbolTable
-> [DataGroup (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
forall a.
StructureTable
-> SymbolTable
-> [DataGroup (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
dataGroups StructureTable
strt SymbolTable
symt (AList DataGroup (Analysis a) -> [DataGroup (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList DataGroup (Analysis a)
groups)
       ]
    [Either TypeError (Type, Expression (Analysis a))]
-> [Either TypeError (Type, Expression (Analysis a))]
-> [Either TypeError (Type, Expression (Analysis a))]
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 <-
         ProgramUnit (Analysis a) -> [Statement (Analysis a)]
forall from to. Biplate from to => from -> [to]
universeBi ProgramUnit (Analysis a)
pu :: [Statement (Analysis a)]
       , Either TypeError (Type, Expression (Analysis a))
res <- StructureTable
-> SymbolTable
-> [Declarator (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
forall a.
StructureTable
-> SymbolTable
-> [Declarator (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
declarators StructureTable
strt SymbolTable
symt ([Declarator (Analysis a)]
 -> [Either TypeError (Type, Expression (Analysis a))])
-> [Declarator (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
forall a b. (a -> b) -> a -> b
$ AList Declarator (Analysis a) -> [Declarator (Analysis a)]
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 = (DataGroup (Analysis a)
 -> [Either TypeError (Type, Expression (Analysis a))])
-> [DataGroup (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DataGroup (Analysis a)
-> [Either TypeError (Type, Expression (Analysis a))]
forall {a}.
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         = AList Expression (Analysis a) -> [Expression (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Expression (Analysis a)
vas
        es :: [Expression (Analysis a)]
es         = AList Expression (Analysis a) -> [Expression (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Expression (Analysis a)
eas
        expandedVs :: [Either TypeError Type]
expandedVs = (Expression (Analysis a) -> [Either TypeError Type])
-> [Expression (Analysis a)] -> [Either TypeError Type]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (StructureTable
-> SymbolTable
-> Expression (Analysis a)
-> [Either TypeError Type]
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
_ = a -> Either a (a, b)
forall a b. a -> Either a b
Left a
err
        g (Right a
ty ) b
e = (a, b) -> Either a (a, b)
forall a b. b -> Either a b
Right (a
ty, b
e)
    in  (Either TypeError Type
 -> Expression (Analysis a)
 -> Either TypeError (Type, Expression (Analysis a)))
-> [Either TypeError Type]
-> [Expression (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Either TypeError Type
-> Expression (Analysis a)
-> Either TypeError (Type, Expression (Analysis a))
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 = (Declarator (Analysis a)
 -> [Either TypeError (Type, Expression (Analysis a))])
-> [Declarator (Analysis a)]
-> [Either TypeError (Type, Expression (Analysis a))]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Declarator (Analysis a)
-> [Either TypeError (Type, Expression (Analysis a))]
forall {a}. Declarator a -> [Either TypeError (Type, Expression a)]
f where
  f :: Declarator a -> [Either TypeError (Type, Expression a)]
f (Declarator a
_ SrcSpan
_ Expression a
v DeclaratorType a
ScalarDecl Maybe (Expression a)
_ (Just Expression a
e)) = Either TypeError (Type, Expression a)
-> [Either TypeError (Type, Expression a)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either TypeError (Type, Expression a)
 -> [Either TypeError (Type, Expression a)])
-> Either TypeError (Type, Expression a)
-> [Either TypeError (Type, Expression a)]
forall a b. (a -> b) -> a -> b
$ (, Expression a
e) (Type -> (Type, Expression a))
-> Either TypeError Type -> Either TypeError (Type, Expression a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructureTable
-> SymbolTable -> Expression a -> Either TypeError Type
forall a.
StructureTable
-> SymbolTable -> Expression a -> Either TypeError Type
typeOf StructureTable
strt SymbolTable
symt Expression a
v
  f d :: Declarator a
d@(Declarator a
_ SrcSpan
_ (ExpValue a
_ SrcSpan
s (ValVariable Name
v)) ArrayDecl{} Maybe (Expression a)
_ (Just (ExpInitialisation a
_ SrcSpan
_ AList Expression a
vals)))
    = case Name -> SymbolTable -> Maybe SymbolTableEntry
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
v SymbolTable
symt of
      Just (SVariable (TArray Type
ty (Just Dimensions
dims)) Location
_) ->
        let tys :: [Type]
tys   = Dimensions -> Type -> [Type]
forall a. Dimensions -> a -> [a]
expandDimensions Dimensions
dims Type
ty
            vals' :: [Expression a]
vals' = AList Expression a -> [Expression a]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Expression a
vals
        in  if [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
tys Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Expression a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Expression a]
vals'
              then
                Either TypeError (Type, Expression a)
-> [Either TypeError (Type, Expression a)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                (Either TypeError (Type, Expression a)
 -> [Either TypeError (Type, Expression a)])
-> ([Name] -> Either TypeError (Type, Expression a))
-> [Name]
-> [Either TypeError (Type, Expression a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeError -> Either TypeError (Type, Expression a)
forall a b. a -> Either a b
Left
                (TypeError -> Either TypeError (Type, Expression a))
-> ([Name] -> TypeError)
-> [Name]
-> Either TypeError (Type, Expression a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcSpan -> Name -> TypeError
typeError SrcSpan
s
                (Name -> TypeError) -> ([Name] -> Name) -> [Name] -> TypeError
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Name] -> Name
forall a. Monoid a => [a] -> a
mconcat
                ([Name] -> [Either TypeError (Type, Expression a)])
-> [Name] -> [Either TypeError (Type, Expression a)]
forall a b. (a -> b) -> a -> b
$ [ Name
"Length of lhs and rhs in declarator do not match: "
                  , Name
": "
                  , Declarator a -> Name
forall a. IndentablePretty a => a -> Name
pprint77l Declarator a
d
                  ]
              else ((Type, Expression a) -> Either TypeError (Type, Expression a))
-> [(Type, Expression a)]
-> [Either TypeError (Type, Expression a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Type, Expression a) -> Either TypeError (Type, Expression a)
forall a b. b -> Either a b
Right ([(Type, Expression a)] -> [Either TypeError (Type, Expression a)])
-> [(Type, Expression a)]
-> [Either TypeError (Type, Expression a)]
forall a b. (a -> b) -> a -> b
$ [Type] -> [Expression a] -> [(Type, Expression a)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Dimensions -> Type -> [Type]
forall a. Dimensions -> a -> [a]
expandDimensions Dimensions
dims Type
ty) ([Expression a] -> [(Type, Expression a)])
-> [Expression a] -> [(Type, Expression a)]
forall a b. (a -> b) -> a -> b
$ AList Expression a -> [Expression a]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Expression a
vals
      Maybe SymbolTableEntry
_ ->
        Either TypeError (Type, Expression a)
-> [Either TypeError (Type, Expression a)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
          (Either TypeError (Type, Expression a)
 -> [Either TypeError (Type, Expression a)])
-> (Name -> Either TypeError (Type, Expression a))
-> Name
-> [Either TypeError (Type, Expression a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  TypeError -> Either TypeError (Type, Expression a)
forall a b. a -> Either a b
Left
          (TypeError -> Either TypeError (Type, Expression a))
-> (Name -> TypeError)
-> Name
-> Either TypeError (Type, Expression a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  SrcSpan -> Name -> TypeError
typeError SrcSpan
s
          (Name -> [Either TypeError (Type, Expression a)])
-> Name -> [Either TypeError (Type, Expression a)]
forall a b. (a -> b) -> a -> b
$  Name
"Unexpected lhs in array declaration at: "
          Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Declarator a -> Name
forall a. IndentablePretty a => a -> Name
pprint77l Declarator a
d
  f Declarator 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 Name -> SymbolTable -> Maybe SymbolTableEntry
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
var SymbolTable
symt of
    Just (SVariable (TArray Type
ty (Just Dimensions
dims)) Location
_) ->
      Dimensions -> Either TypeError Type -> [Either TypeError Type]
forall a. Dimensions -> a -> [a]
expandDimensions Dimensions
dims (Type -> Either TypeError Type
forall a b. b -> Either a b
Right Type
ty)
    Just (SVariable Type
ty Location
_) -> [Type -> Either TypeError Type
forall a b. b -> Either a b
Right Type
ty]
    Maybe SymbolTableEntry
_ ->
      Either TypeError Type -> [Either TypeError Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        (Either TypeError Type -> [Either TypeError Type])
-> (Name -> Either TypeError Type)
-> Name
-> [Either TypeError Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  TypeError -> Either TypeError Type
forall a b. a -> Either a b
Left
        (TypeError -> Either TypeError Type)
-> (Name -> TypeError) -> Name -> Either TypeError Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  SrcSpan -> Name -> TypeError
typeError (Expression (Analysis a) -> SrcSpan
forall a. Spanned a => a -> SrcSpan
getSpan Expression (Analysis a)
e)
        (Name -> [Either TypeError Type])
-> Name -> [Either TypeError Type]
forall a b. (a -> b) -> a -> b
$  Name
"Got unexpected lhs type: "
        Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Expression (Analysis a) -> Name
forall a. IndentablePretty a => a -> Name
pprint77l Expression (Analysis a)
e
  Expression (Analysis a)
_ -> Either TypeError Type -> [Either TypeError Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either TypeError Type -> [Either TypeError Type])
-> Either TypeError Type -> [Either TypeError Type]
forall a b. (a -> b) -> a -> b
$ StructureTable
-> SymbolTable -> Expression (Analysis a) -> Either TypeError Type
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 :: Dimensions -> a -> [a]
expandDimensions :: forall a. Dimensions -> a -> [a]
expandDimensions Dimensions
dims =
  Int -> a -> [a]
forall a. Int -> a -> [a]
replicate ((Int -> (Int, Int) -> Int) -> Int -> Dimensions -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
acc (Int
x, Int
y) -> Int -> Int
forall a. Num a => a -> a
abs (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
acc) Int
1 Dimensions
dims)