-- GRIN-like backend for Yhc Core.
-- Seed the Heap Map from Core annotations.

module Yhc.Core.GRIN.HeapSeed where

import Yhc.Core.Extra
import Yhc.Core.GRIN.Type
import Yhc.Core.GRIN.SubstVars
import Yhc.Core.GRIN.HeapPointsTo
import qualified Data.Map as M
import Data.List
import Data.Char

import Text.ParserCombinators.Parsec
import Text.ParserCombinators.Parsec.Token
import Text.ParserCombinators.Parsec.Language

-- Given a linked Yhc Core and Core Annotations, obtain a seed for the Heap Map. 

coreHeapSeed :: Core -> CoreAnnotations -> HeapMap

coreHeapSeed core anno =
  let prims = filter isCorePrim (coreFuncs core)
      tsigs = map (getTSIG . getTypeSig anno) prims
      psigs = map getRetType tsigs
      pnames = map coreFuncName prims
      tsmap = M.fromList $ zipWith (\a b -> (HFunc a, HasType b)) pnames tsigs
      rtmap = M.unions $ zipWith (tsig2pt core) pnames psigs
  in  M.intersectionWith (:) tsmap rtmap

-- Obtain possible values from a type signature. Types normally expected
-- to be returned from primitives or passed from the execution platform
-- to Haskell code are supported, that is:
--
--  * Integral types (Integer, Float, Char): represented as generalized values
--  * Algebraic types with zero-ary constructors: represented as union of GTags
--  * Algebraic types with one constructor of any arity (e. g. tuples): represented
--      as single GTagged values with arguments converted
--  * List types with some special arrangements to reflect recursive structure of
--      the type
--
--  The rest (like Maybe or Either) will only have GTany as possible values of arguments.

tsig2pt :: Core -> GName -> TSIG -> HeapMap

tsig2pt core gn (Terr _) = M.empty

tsig2pt core gn (Tvar _) = M.singleton (HFunc gn) [HasValue $ GVal GTAny]

tsig2pt core gn (Tcon "Int" []) = M.singleton (HFunc gn) [HasValue $ GVal GTInteger]
tsig2pt core gn (Tcon "Integer" []) = M.singleton (HFunc gn) [HasValue $ GVal GTInteger]
tsig2pt core gn (Tcon "Float" []) = M.singleton (HFunc gn) [HasValue $ GVal GTDouble]
tsig2pt core gn (Tcon "Double" []) = M.singleton (HFunc gn) [HasValue $ GVal GTDouble]
tsig2pt core gn (Tcon "Char" []) = M.singleton (HFunc gn) [HasValue $ GVal GTChar]
tsig2pt core gn (Tcon "String" []) = tsig2pt core gn (Tcon "Prelude;[]" [Tcon "Char" []])

tsig2pt core gn (Tcon con args) = 
  let cd = case coreDataMaybe core con of
        Just ccd -> ccd
        Nothing -> CoreData {
          coreDataName = con
         ,coreDataTypes = []
         ,coreDataCtors = [CoreCtor {
           coreCtorName = con
          ,coreCtorFields = replicate (length args) ("*", Nothing)}]}
      islist = con == "Prelude;[]" && nargs == 1
      ctors = coreDataCtors cd
      nctors = length ctors
      nargs = length args
      ctarity = length . coreCtorFields
      maxarity = maximum (map ctarity ctors)
      nargv = nargs + if islist then 1 else 0
      argns = map ((\s -> "n'" ++ con ++ "'" ++ gn ++ "'" ++ s) . show) (take nargv [1 .. ])
      argmap argn argt = fixmap argn $ tsig2pt core argn argt
      fixmap n m = case M.lookup (HFunc n) m of
        Nothing -> m
        Just v -> M.delete (HFunc n) m `M.union` M.singleton (HVar n) v
      argmaps = zipWith argmap argns args
      subctor ctor = let tag = mkTagName $ coreCtorName ctor in case ctarity ctor of
        0 -> HasValue $ GTag tag
        n -> HasValue (GTagged tag (take n (map GVar argns)))
      anyctor ctor = let tag = mkTagName $ coreCtorName ctor in case ctarity ctor of
        0 -> HasValue $ GTag tag
        n -> HasValue (GTagged tag (replicate n GTAny))
      funmap = M.singleton (HFunc gn) (map subctor ctors)
      anymap = M.singleton (HFunc gn) (map anyctor ctors)
      retmap = M.singleton (HVar (head $ reverse argns))
                           [ReturnOf gn]
  in  case (nctors, maxarity, islist) of
        (_, 0, False) -> funmap
        (1, _, False) -> M.unions (funmap : argmaps)
        (_, _, True)  -> M.unions (funmap : retmap : argmaps)
        (_, _, _)     -> anymap


-- Obtain a return type from type signature.

getRetType :: TSIG -> TSIG

getRetType (Tapp f a) = getRetType a
getRetType z = z

-- Obtain argument types from type signature.

getArgTypes :: TSIG -> [TSIG]

getArgTypes (Tapp f a) = f : getArgTypes a
getArgTypes _ = []

-- Given a Core function and annotations (possibly combined from several sources),
-- retrieve function's type signature if available. If annotation is not available,
-- return an empty string.

getTypeSig :: CoreAnnotations -> CoreFunc -> String

getTypeSig anno cf =
  case getAnnotation cf "Type" anno of
    Just (CoreTypeSig s) -> s
    _ -> ""

-- A simple parser for Haskell type expressions.
-- Parsing is done in accordance with the following productions (Haskell report, 4.1.2):
--
-- type   ->  btype [-> type]           (function type)
-- btype  ->  [btype] atype             (type application)
-- atype  ->  gtycon
--        |   tyvar
--        |   ( type1 , ... , typek ) 	(tuple type, k>=2)
--        |   [ type ]                  (list type)
--        |   ( type )                  (parenthesised constructor)
-- gtycon ->  qtycon
--        |   ()                        (unit type)
--        |   []                        (list constructor)
--        |   (->)                      (function constructor)
--        |   (,{,})                    (tupling constructors) 
--
-- Since this parser is intended to only parse type signatures of primitives,
-- many of real Haskell type system features are not implemented here.
--
-- Type constructors (tags) are expested to start with a capital letter and contain any
-- non-whitespace characters. Type variables are expected to start with a lowercase letter
-- and contain any non-whitespace characters. Function type constructors (->) are treated as
-- separators.

-- Tokenizer

rsrvd = [",", "->", "(", ")", "[", "]"]

idLetter x = not (isSpace x) && x `notElem` concat rsrvd

tok = makeTokenParser emptyDef {
  identStart = upper <|> lower
 ,identLetter = satisfy idLetter
 ,reservedNames = rsrvd
}

-- Run the parser

getTSIG :: String -> TSIG

getTSIG s = case parse parseTSIG "" s of
  Left err -> Terr (show err)
  Right ts -> ts

parseTSIG :: Parser TSIG

parseTSIG = do
  t <- oneTSIG
  ts <- many (symbol tok "->" >> parseTSIG)
  case ts of
    [] -> return t
    (ts:_) -> return $ Tapp t ts

oneTSIG :: Parser TSIG

oneTSIG = var <|> con <|> fun <|> tuple <|> list

var = try $ do
  v <- lower
  ar <- many $ satisfy idLetter
  skipMany space
  return $ Tvar (v:ar)

con = try $ do
  t <- tag
  args <- many oneTSIG
  return $ Tcon t args

tag = do
  t <- upper
  ag <- many $ satisfy idLetter
  skipMany space
  return (t:ag)

fun = try $ do
  symbol tok "("
  t <- parseTSIG
  symbol tok ")"
  return t

tuple = try $ do
  symbol tok "("
  ts <- sepBy oneTSIG (symbol tok ",")
  symbol tok ")"
  case ts of
    []  -> return $ Tcon "Prelude;()" []
    [t] -> return t
    tts -> return $ Tcon ("Prelude;(" ++ replicate (length tts - 1) ',' ++ ")") tts

list = try $ do
  symbol tok "[" 
  t <- oneTSIG
  symbol tok "]"
  return $ Tcon "[]" [t]