{-# 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) -> [StructureTableEntry]
itemToEntry SymbolTable
st (StructFields Analysis a
_ SrcSpan
_ TypeSpec (Analysis a)
ty Maybe (AList Attribute (Analysis a))
_ AList Declarator (Analysis a)
decls) =
  forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (forall a.
SymbolTable
-> TypeSpec (Analysis a)
-> Declarator (Analysis a)
-> Maybe StructureTableEntry
handleDeclarator SymbolTable
st TypeSpec (Analysis a)
ty) (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) = [[[StructureTableEntry]] -> StructureTableEntry
UnionEntry forall a b. (a -> b) -> a -> b
$ forall a.
SymbolTable -> UnionMap (Analysis a) -> [StructureTableEntry]
handleUnion SymbolTable
st forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 = forall a. SymbolTable -> TypeSpec (Analysis a) -> Type
typeSpecToScalarType SymbolTable
st TypeSpec (Analysis a)
ty
  in  forall a. Expression a -> Maybe String
expToName Expression (Analysis a)
expr forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \String
name -> forall a. a -> Maybe a
Just 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 = forall a.
SymbolTable
-> [DimensionDeclarator (Analysis a)]
-> TypeSpec (Analysis a)
-> Type
typeSpecToArrayType SymbolTable
st (forall (t :: * -> *) a. AList t a -> [t a]
aStrip AList DimensionDeclarator (Analysis a)
dims) TypeSpec (Analysis a)
ty
  in  forall a. Expression a -> Maybe String
expToName Expression (Analysis a)
expr forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \String
name -> forall a. a -> Maybe a
Just 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) -> [StructureTableEntry]
handleUnion SymbolTable
st (UnionMap Analysis a
_ SrcSpan
_ AList StructureItem (Analysis a)
si) = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a.
SymbolTable -> StructureItem (Analysis a) -> [StructureTableEntry]
itemToEntry SymbolTable
st) forall a b. (a -> b) -> a -> b
$ 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)) = forall a. a -> Maybe a
Just String
name
expToName Expression a
_ = 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 =
  forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' StructureTable
-> String -> [StructureItem (Analysis a)] -> StructureTable
handler forall k a. Map k a
M.empty
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
    forall a b. (a -> b) -> a -> b
$ [ (String
n, 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) <- forall a. Data a => ProgramUnit a -> [Statement a]
allPUS ProgramUnit (Analysis a)
pu ]
 where
  handler :: StructureTable
-> String -> [StructureItem (Analysis a)] -> StructureTable
handler StructureTable
structTable String
name [StructureItem (Analysis a)]
entry =
    forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert String
name (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a.
SymbolTable -> StructureItem (Analysis a) -> [StructureTableEntry]
itemToEntry SymbolTable
st) [StructureItem (Analysis a)]
entry) StructureTable
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 forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
tyName StructureTable
structTable of
    Just [StructureTableEntry]
struct -> String -> [StructureTableEntry] -> Either TypeError Type
hasEntry String
ref [StructureTableEntry]
struct
    Maybe [StructureTableEntry]
Nothing     -> forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> TypeError
UnknownField forall a b. (a -> b) -> a -> b
$ String
ref forall a. Semigroup a => a -> a -> a
<> String
" not a field of " forall a. Semigroup a => a -> a -> a
<> String
tyName
  Type
ty' -> forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> TypeError
UnknownField forall a b. (a -> b) -> a -> b
$ String
"No fields for data type (" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Type
ty' 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 -> [StructureTableEntry] -> Either TypeError Type
hasEntry String
name [StructureTableEntry]
struct =
  let unionStructs :: [StructureTableEntry]
unionStructs = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[StructureTableEntry]]
structs | (UnionEntry [[StructureTableEntry]]
structs) <- [StructureTableEntry]
struct ]
  in
    case
      [ Type
ty | (FieldEntry String
fname Type
ty) <- [StructureTableEntry]
struct forall a. Semigroup a => a -> a -> a
<> [StructureTableEntry]
unionStructs, String
fname forall a. Eq a => a -> a -> Bool
== String
name ]
    of
      [Type
ty] -> forall a b. b -> Either a b
Right Type
ty
      []   -> forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> TypeError
UnknownField forall a b. (a -> b) -> a -> b
$ String
name forall a. Semigroup a => a -> a -> a
<> String
" is not a field"
      [Type]
_ ->
        forall a b. a -> Either a b
Left
          forall b c a. (b -> c) -> (a -> b) -> a -> c
.  String -> TypeError
UnknownField
          forall a b. (a -> b) -> a -> b
$  String
"Field "
          forall a. Semigroup a => a -> a -> a
<> String
name
          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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {a}.
Data a =>
ProgramStructureTables
-> ProgramUnit (Analysis a) -> ProgramStructureTables
handler forall k a. Map k a
M.empty forall a b. (a -> b) -> a -> b
$ 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 = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall a. ProgramUnit (Analysis a) -> ProgramUnitName
puName ProgramUnit (Analysis a)
pu) (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 = forall a. Data a => ProgramUnit (Analysis a) -> SymbolTable
collectSymbols ProgramUnit (Analysis a)
pu in forall a.
Data a =>
SymbolTable -> ProgramUnit (Analysis a) -> StructureTable
collectStructures SymbolTable
st ProgramUnit (Analysis a)
pu