{-# LANGUAGE ScopedTypeVariables #-}

module Language.Fortran.Vars.StructureTable
  ( StructureTable
  , StructureTableEntry(..)
  , collectStructures
  , lookupField
  , hasEntry
  , programStructureTables
  )
where

import           Data.Data                      ( Data )
import           Data.Maybe                     ( mapMaybe )
import qualified Data.Map                      as M
import           Data.List                      ( foldl' )

import           Language.Fortran.Analysis      ( Analysis
                                                , puName
                                                )
import           Language.Fortran.AST           ( Statement(..)
                                                , StructureItem(..)
                                                , UnionMap(..)
                                                , Expression(..)
                                                , Value(..)
                                                , ProgramUnit(..)
                                                , ProgramFile
                                                , TypeSpec(..)
                                                , Declarator(..)
                                                , DeclaratorType(..)
                                                , aStrip
                                                )
import           Language.Fortran.Extras
                                                ( allPUS
                                                , allPU
                                                )

import           Language.Fortran.Vars.SymbolTable
                                                ( collectSymbols )
import           Language.Fortran.Vars.Types
                                                ( SymbolTable
                                                , StructureTableEntry(..)
                                                , Structure
                                                , StructureTable
                                                , ProgramStructureTables
                                                , Type(..)
                                                , SemType(..)
                                                , TypeError(..)
                                                )
import           Language.Fortran.Vars.Utils
                                                ( typeSpecToArrayType
                                                , typeSpecToScalarType
                                                )

-- | Given a `SymbolTable` transform `StructureItem`s found in the AST into a list of
-- `StructureTableEntry`s
itemToEntry
  :: SymbolTable -> StructureItem (Analysis a) -> [StructureTableEntry]
itemToEntry :: forall a. SymbolTable -> StructureItem (Analysis a) -> Structure
itemToEntry SymbolTable
st (StructFields Analysis a
_ SrcSpan
_ TypeSpec (Analysis a)
ty Maybe (AList Attribute (Analysis a))
_ AList Declarator (Analysis a)
decls) =
  (Declarator (Analysis a) -> Maybe StructureTableEntry)
-> [Declarator (Analysis a)] -> Structure
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (SymbolTable
-> TypeSpec (Analysis a)
-> Declarator (Analysis a)
-> Maybe StructureTableEntry
forall a.
SymbolTable
-> TypeSpec (Analysis a)
-> Declarator (Analysis a)
-> Maybe StructureTableEntry
handleDeclarator SymbolTable
st TypeSpec (Analysis a)
ty) (AList Declarator (Analysis a) -> [Declarator (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList Declarator (Analysis a)
decls)
itemToEntry SymbolTable
st (StructUnion Analysis a
_ SrcSpan
_ AList UnionMap (Analysis a)
l) = [[Structure] -> StructureTableEntry
UnionEntry ([Structure] -> StructureTableEntry)
-> [Structure] -> StructureTableEntry
forall a b. (a -> b) -> a -> b
$ SymbolTable -> UnionMap (Analysis a) -> Structure
forall a. SymbolTable -> UnionMap (Analysis a) -> Structure
handleUnion SymbolTable
st (UnionMap (Analysis a) -> Structure)
-> [UnionMap (Analysis a)] -> [Structure]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AList UnionMap (Analysis a) -> [UnionMap (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList UnionMap (Analysis a)
l]
itemToEntry SymbolTable
_  StructStructure{}   = []

-- TODO take into account length, should override default typespecs
-- | Given the 'TypeSpec' and 'Declarator' found in a field entry create a
--   'StructureTableEntry'
handleDeclarator
  :: SymbolTable
  -> TypeSpec (Analysis a)
  -> Declarator (Analysis a)
  -> Maybe StructureTableEntry
handleDeclarator :: forall a.
SymbolTable
-> TypeSpec (Analysis a)
-> Declarator (Analysis a)
-> Maybe StructureTableEntry
handleDeclarator SymbolTable
st TypeSpec (Analysis a)
ty (Declarator Analysis a
_ SrcSpan
_ Expression (Analysis a)
expr DeclaratorType (Analysis a)
ScalarDecl Maybe (Expression (Analysis a))
_ Maybe (Expression (Analysis a))
_) =
  let scalarTy :: Type
scalarTy = SymbolTable -> TypeSpec (Analysis a) -> Type
forall a. SymbolTable -> TypeSpec (Analysis a) -> Type
typeSpecToScalarType SymbolTable
st TypeSpec (Analysis a)
ty
  in  Expression (Analysis a) -> Maybe String
forall a. Expression a -> Maybe String
expToName Expression (Analysis a)
expr Maybe String
-> (String -> Maybe StructureTableEntry)
-> Maybe StructureTableEntry
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \String
name -> StructureTableEntry -> Maybe StructureTableEntry
forall a. a -> Maybe a
Just (StructureTableEntry -> Maybe StructureTableEntry)
-> StructureTableEntry -> Maybe StructureTableEntry
forall a b. (a -> b) -> a -> b
$ String -> Type -> StructureTableEntry
FieldEntry String
name Type
scalarTy
handleDeclarator SymbolTable
st TypeSpec (Analysis a)
ty (Declarator Analysis a
_ SrcSpan
_ Expression (Analysis a)
expr (ArrayDecl AList DimensionDeclarator (Analysis a)
dims) Maybe (Expression (Analysis a))
_ Maybe (Expression (Analysis a))
_) =
  let arrayty :: Type
arrayty = SymbolTable
-> [DimensionDeclarator (Analysis a)]
-> TypeSpec (Analysis a)
-> Type
forall a.
SymbolTable
-> [DimensionDeclarator (Analysis a)]
-> TypeSpec (Analysis a)
-> Type
typeSpecToArrayType SymbolTable
st (AList DimensionDeclarator (Analysis a)
-> [DimensionDeclarator (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList DimensionDeclarator (Analysis a)
dims) TypeSpec (Analysis a)
ty
  in  Expression (Analysis a) -> Maybe String
forall a. Expression a -> Maybe String
expToName Expression (Analysis a)
expr Maybe String
-> (String -> Maybe StructureTableEntry)
-> Maybe StructureTableEntry
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \String
name -> StructureTableEntry -> Maybe StructureTableEntry
forall a. a -> Maybe a
Just (StructureTableEntry -> Maybe StructureTableEntry)
-> StructureTableEntry -> Maybe StructureTableEntry
forall a b. (a -> b) -> a -> b
$ String -> Type -> StructureTableEntry
FieldEntry String
name Type
arrayty

-- | Transform a `UnionMap` in an AST to `StructureTableEntry`s
handleUnion :: SymbolTable -> UnionMap (Analysis a) -> [StructureTableEntry]
handleUnion :: forall a. SymbolTable -> UnionMap (Analysis a) -> Structure
handleUnion SymbolTable
st (UnionMap Analysis a
_ SrcSpan
_ AList StructureItem (Analysis a)
si) = (StructureItem (Analysis a) -> Structure)
-> [StructureItem (Analysis a)] -> Structure
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (SymbolTable -> StructureItem (Analysis a) -> Structure
forall a. SymbolTable -> StructureItem (Analysis a) -> Structure
itemToEntry SymbolTable
st) ([StructureItem (Analysis a)] -> Structure)
-> [StructureItem (Analysis a)] -> Structure
forall a b. (a -> b) -> a -> b
$ AList StructureItem (Analysis a) -> [StructureItem (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList StructureItem (Analysis a)
si

-- | Given an `Expression` maybe get the name
expToName :: Expression a -> Maybe String
expToName :: forall a. Expression a -> Maybe String
expToName (ExpValue a
_ SrcSpan
_ (ValVariable String
name)) = String -> Maybe String
forall a. a -> Maybe a
Just String
name
expToName Expression a
_ = Maybe String
forall a. Maybe a
Nothing

-- | Collect structures defined in a `ProgramUnit` and return a `StructureTable`
collectStructures
  :: Data a => SymbolTable -> ProgramUnit (Analysis a) -> StructureTable
collectStructures :: forall a.
Data a =>
SymbolTable -> ProgramUnit (Analysis a) -> StructureTable
collectStructures SymbolTable
st ProgramUnit (Analysis a)
pu =
  (StructureTable
 -> String -> [StructureItem (Analysis a)] -> StructureTable)
-> StructureTable
-> Map String [StructureItem (Analysis a)]
-> StructureTable
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' StructureTable
-> String -> [StructureItem (Analysis a)] -> StructureTable
forall {k} {t :: * -> *} {a}.
(Ord k, Foldable t) =>
Map k Structure
-> k -> t (StructureItem (Analysis a)) -> Map k Structure
handler StructureTable
forall k a. Map k a
M.empty
    (Map String [StructureItem (Analysis a)] -> StructureTable)
-> ([(String, [StructureItem (Analysis a)])]
    -> Map String [StructureItem (Analysis a)])
-> [(String, [StructureItem (Analysis a)])]
-> StructureTable
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(String, [StructureItem (Analysis a)])]
-> Map String [StructureItem (Analysis a)]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
    ([(String, [StructureItem (Analysis a)])] -> StructureTable)
-> [(String, [StructureItem (Analysis a)])] -> StructureTable
forall a b. (a -> b) -> a -> b
$ [ (String
n, AList StructureItem (Analysis a) -> [StructureItem (Analysis a)]
forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList StructureItem (Analysis a)
s) | (StStructure Analysis a
_ SrcSpan
_ (Just String
n) AList StructureItem (Analysis a)
s) <- ProgramUnit (Analysis a) -> [Statement (Analysis a)]
forall a. Data a => ProgramUnit a -> [Statement a]
allPUS ProgramUnit (Analysis a)
pu ]
 where
  handler :: Map k Structure
-> k -> t (StructureItem (Analysis a)) -> Map k Structure
handler Map k Structure
structTable k
name t (StructureItem (Analysis a))
entry =
    k -> Structure -> Map k Structure -> Map k Structure
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
name ((StructureItem (Analysis a) -> Structure)
-> t (StructureItem (Analysis a)) -> Structure
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (SymbolTable -> StructureItem (Analysis a) -> Structure
forall a. SymbolTable -> StructureItem (Analysis a) -> Structure
itemToEntry SymbolTable
st) t (StructureItem (Analysis a))
entry) Map k Structure
structTable

-- | Look up reference on a type to return another `Type`
lookupField :: StructureTable -> Type -> String -> Either TypeError Type
lookupField :: StructureTable -> Type -> String -> Either TypeError Type
lookupField StructureTable
structTable Type
ty String
ref = case Type
ty of
  TCustom String
tyName -> case String -> StructureTable -> Maybe Structure
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
tyName StructureTable
structTable of
    Just Structure
struct -> String -> Structure -> Either TypeError Type
hasEntry String
ref Structure
struct
    Maybe Structure
Nothing     -> TypeError -> Either TypeError Type
forall a b. a -> Either a b
Left (TypeError -> Either TypeError Type)
-> (String -> TypeError) -> String -> Either TypeError Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> TypeError
UnknownField (String -> Either TypeError Type)
-> String -> Either TypeError Type
forall a b. (a -> b) -> a -> b
$ String
ref String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" not a field of " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
tyName
  Type
ty' -> TypeError -> Either TypeError Type
forall a b. a -> Either a b
Left (TypeError -> Either TypeError Type)
-> (String -> TypeError) -> String -> Either TypeError Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> TypeError
UnknownField (String -> Either TypeError Type)
-> String -> Either TypeError Type
forall a b. (a -> b) -> a -> b
$ String
"No fields for data type (" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Type -> String
forall a. Show a => a -> String
show Type
ty' String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
")"

-- | Given a name, check that a `Structure` contains it once and return its
-- corresponding `Type`
hasEntry :: String -> Structure -> Either TypeError Type
hasEntry :: String -> Structure -> Either TypeError Type
hasEntry String
name Structure
struct =
  let unionStructs :: Structure
unionStructs = [Structure] -> Structure
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ [Structure] -> Structure
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Structure]
structs | (UnionEntry [Structure]
structs) <- Structure
struct ]
  in
    case
      [ Type
ty | (FieldEntry String
fname Type
ty) <- Structure
struct Structure -> Structure -> Structure
forall a. Semigroup a => a -> a -> a
<> Structure
unionStructs, String
fname String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
name ]
    of
      [Type
ty] -> Type -> Either TypeError Type
forall a b. b -> Either a b
Right Type
ty
      []   -> TypeError -> Either TypeError Type
forall a b. a -> Either a b
Left (TypeError -> Either TypeError Type)
-> (String -> TypeError) -> String -> Either TypeError Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> TypeError
UnknownField (String -> Either TypeError Type)
-> String -> Either TypeError Type
forall a b. (a -> b) -> a -> b
$ String
name String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" is not a field"
      [Type]
_ ->
        TypeError -> Either TypeError Type
forall a b. a -> Either a b
Left
          (TypeError -> Either TypeError Type)
-> (String -> TypeError) -> String -> Either TypeError Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  String -> TypeError
UnknownField
          (String -> Either TypeError Type)
-> String -> Either TypeError Type
forall a b. (a -> b) -> a -> b
$  String
"Field "
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
name
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" define multiple times for structure"

-- | Given a 'ProgramFile', generate a 'StructureTable' for each 'ProgramUnit'.
-- This can be used to check types in data reference expressions
programStructureTables
  :: Data a => ProgramFile (Analysis a) -> ProgramStructureTables
programStructureTables :: forall a.
Data a =>
ProgramFile (Analysis a) -> ProgramStructureTables
programStructureTables ProgramFile (Analysis a)
pf = (ProgramStructureTables
 -> ProgramUnit (Analysis a) -> ProgramStructureTables)
-> ProgramStructureTables
-> [ProgramUnit (Analysis a)]
-> ProgramStructureTables
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ProgramStructureTables
-> ProgramUnit (Analysis a) -> ProgramStructureTables
forall {a}.
Data a =>
ProgramStructureTables
-> ProgramUnit (Analysis a) -> ProgramStructureTables
handler ProgramStructureTables
forall k a. Map k a
M.empty ([ProgramUnit (Analysis a)] -> ProgramStructureTables)
-> [ProgramUnit (Analysis a)] -> ProgramStructureTables
forall a b. (a -> b) -> a -> b
$ ProgramFile (Analysis a) -> [ProgramUnit (Analysis a)]
forall a. Data a => ProgramFile a -> [ProgramUnit a]
allPU ProgramFile (Analysis a)
pf
 where
  handler :: ProgramStructureTables
-> ProgramUnit (Analysis a) -> ProgramStructureTables
handler ProgramStructureTables
m ProgramUnit (Analysis a)
pu = ProgramUnitName
-> StructureTable
-> ProgramStructureTables
-> ProgramStructureTables
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (ProgramUnit (Analysis a) -> ProgramUnitName
forall a. ProgramUnit (Analysis a) -> ProgramUnitName
puName ProgramUnit (Analysis a)
pu) (ProgramUnit (Analysis a) -> StructureTable
forall {a}. Data a => ProgramUnit (Analysis a) -> StructureTable
puStructure ProgramUnit (Analysis a)
pu) ProgramStructureTables
m
  puStructure :: ProgramUnit (Analysis a) -> StructureTable
puStructure ProgramUnit (Analysis a)
pu = let st :: SymbolTable
st = ProgramUnit (Analysis a) -> SymbolTable
forall a. Data a => ProgramUnit (Analysis a) -> SymbolTable
collectSymbols ProgramUnit (Analysis a)
pu in SymbolTable -> ProgramUnit (Analysis a) -> StructureTable
forall a.
Data a =>
SymbolTable -> ProgramUnit (Analysis a) -> StructureTable
collectStructures SymbolTable
st ProgramUnit (Analysis a)
pu