-- |
-- Module      :  Cryptol.TypeCheck.Solver.Selector
-- Copyright   :  (c) 2013-2016 Galois, Inc.
-- License     :  BSD3
-- Maintainer  :  cryptol@galois.com
-- Stability   :  provisional
-- Portability :  portable
{-# LANGUAGE PatternGuards, Safe #-}
module Cryptol.TypeCheck.Solver.Selector (tryHasGoal) where

import Cryptol.TypeCheck.AST
import Cryptol.TypeCheck.InferTypes
import Cryptol.TypeCheck.Monad( InferM, unify, newGoals, lookupNewtype
                              , newType, applySubst, solveHasGoal
                              , newParamName
                              )
import Cryptol.TypeCheck.Subst(listSubst,apSubst)
import Cryptol.Utils.Ident (Ident, packIdent)
import Cryptol.Utils.Panic(panic)

import Control.Monad(forM,guard)


recordType :: [Ident] -> InferM Type
recordType labels =
  do fields <- forM labels $ \l ->
        do t <- newType (TypeOfRecordField l) KType
           return (l,t)
     return (TRec fields)

tupleType :: Int -> InferM Type
tupleType n =
  do fields <- mapM (\x -> newType (TypeOfTupleField x) KType)
                    [ 0 .. (n-1) ]
     return (tTuple fields)

listType :: Int -> InferM Type
listType n =
  do elems <- newType TypeOfSeqElement KType
     return (tSeq (tNum n) elems)


improveSelector :: Selector -> Type -> InferM Bool
improveSelector sel outerT =
  case sel of
    RecordSel _ mb -> cvt recordType mb
    TupleSel  _ mb -> cvt tupleType  mb
    ListSel   _ mb -> cvt listType   mb
  where
  cvt _ Nothing   = return False
  cvt f (Just a)  = do ty <- f a
                       newGoals CtExactType =<< unify ty outerT
                       newT <- applySubst outerT
                       return (newT /= outerT)


{- | Compute the type of a field based on the selector.
The given type should be "zonked" (i.e., substitution was applied to it),
and (outermost) type synonyms have been expanded.
-}
solveSelector :: Selector -> Type -> InferM (Maybe Type)
solveSelector sel outerT =
  case (sel, outerT) of

    (RecordSel l _, ty) ->
      case ty of
        TRec fs  -> return (lookup l fs)
        TCon (TC TCSeq) [len,el] -> liftSeq len el
        TCon (TC TCFun) [t1,t2]  -> liftFun t1 t2
        TCon (TC (TCNewtype (UserTC x _))) ts ->
          do mb <- lookupNewtype x
             case mb of
               Nothing -> return Nothing
               Just nt ->
                 case lookup l (ntFields nt) of
                   Nothing -> return Nothing
                   Just t  ->
                     do let su = listSubst (zip (map tpVar (ntParams nt)) ts)
                        newGoals (CtPartialTypeFun $ UserTyFun x)
                          $ apSubst su $ ntConstraints nt
                        return $ Just $ apSubst su t
        _ -> return Nothing

    (TupleSel n _, ty) ->
      case ty of

        TCon (TC (TCTuple m)) ts ->
          return $ do guard (0 <= n && n < m)
                      return $ ts !! n

        TCon (TC TCSeq) [len,el] -> liftSeq len el
        TCon (TC TCFun) [t1,t2]  -> liftFun t1 t2

        _ -> return Nothing


    (ListSel n _, TCon (TC TCSeq) [l,t])
      | n < 2 -> return (Just t)
      | otherwise ->
       do newGoals CtSelector [ l >== tNum (n - 1) ]
          return (Just t)

    _ -> return Nothing

  where
  liftSeq len el =
    do mb <- solveSelector sel (tNoUser el)
       return $ do el' <- mb
                   return (TCon (TC TCSeq) [len,el'])

  liftFun t1 t2 =
    do mb <- solveSelector sel (tNoUser t2)
       return $ do t2' <- mb
                   return (TCon (TC TCFun) [t1,t2'])


-- | Solve has-constraints.
tryHasGoal :: HasGoal -> InferM (Bool, Bool) -- ^ changes, solved
tryHasGoal has
  | TCon (PC (PHas sel)) [ th, ft ] <- goal (hasGoal has) =
    do imped     <- improveSelector sel th
       outerT    <- tNoUser `fmap` applySubst th
       mbInnerT  <- solveSelector sel outerT
       case mbInnerT of
         Nothing -> return (imped, False)
         Just innerT ->
           do newGoals CtExactType =<< unify innerT ft
              oT <- applySubst outerT
              iT <- applySubst innerT
              selFrom <- mkSel sel oT iT
              solveHasGoal (hasName has) selFrom
              return (True, True)

  | otherwise = panic "hasGoalSolved"
                  [ "Unexpected selector proposition:"
                  , show (hasGoal has)
                  ]

{- | Generator an appropriate selector, once the "Has" constraint
has been discharged.  The resulting selectors should always work
on their corresponding types (i.e., tuple selectros only select from tuples).
This function generates the code for lifting tuple/record selectors to sequences
and functions.

Assumes types are zonked. -}
mkSel :: Selector -> Type -> Type -> InferM (Expr -> Expr)
mkSel s outerT innerT =
  case tNoUser outerT of
    TCon (TC TCSeq) [len,el]
      | TupleSel {} <- s  -> liftSeq len el
      | RecordSel {} <- s -> liftSeq len el

    TCon (TC TCFun) [t1,t2]
      | TupleSel {} <- s -> liftFun t1 t2
      | RecordSel {} <- s -> liftFun t1 t2

    _ -> return (\e -> ESel e s)

  where
  liftSeq len el =
    do x <- newParamName (packIdent "x")
       case tNoUser innerT of
         TCon _ [_,eli] ->
           do selFrom <- mkSel s el eli
              return $ \e -> EComp len eli (selFrom (EVar x))
                                                        [[ From x len el e ]]
         _ -> panic "mkSel" [ "Unexpected inner seq type.", show innerT ]

  liftFun t1 t2 =
    do x <- newParamName (packIdent "x")
       case tNoUser innerT of
         TCon _ [_,inT] ->
           do selFrom <- mkSel s t2 inT
              return $ \e -> EAbs x t1 (selFrom (EApp e (EVar x)))
         _ -> panic "mkSel" [ "Unexpected inner fun type", show innerT ]