ExpandPropGuards.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}

-- |
-- Module      :  Cryptol.Parser.PropGuards
-- Copyright   :  (c) 2022 Galois, Inc.
-- License     :  BSD3
-- Maintainer  :  cryptol@galois.com
-- Stability   :  provisional
-- Portability :  portable
--
-- Expands PropGuards into a top-level definition for each case, and rewrites
-- the body of each case to be an appropriate call to the respectively generated
-- function.
module Cryptol.Parser.ExpandPropGuards where

import Control.DeepSeq
import Cryptol.Parser.AST
import Cryptol.Utils.PP
import Cryptol.Utils.Panic (panic)
import Data.Foldable (traverse_)
import Data.Text (pack)
import GHC.Generics (Generic)

-- | Monad
type ExpandPropGuardsM a = Either Error a

runExpandPropGuardsM :: ExpandPropGuardsM a -> Either Error a
runExpandPropGuardsM m = m

-- | Errors that can happen while expanding constraint guards.
data Error = NoSignature (Located PName)
             -- ^ A declaration using constraints guard lacks an explicit type
             -- signature, which is a requirement.
           | NestedConstraintGuard (Located PName)
             -- ^ Constraint guards appear somewhere other than the top level,
             -- which is not allowed.
  deriving (Show, Generic, NFData)

instance PP Error where
  ppPrec _ err = case err of
    NoSignature x ->
      hang
        (text "At" <+> pp (srcRange x) <.> colon)
        2
        (vcat [ text "Declaration" <+> backticks (pp (thing x)) <+>
                text "lacks a type signature."
              , text "Declarations using constraint guards require an" <+>
                text "explicit type signature."
              ])
    NestedConstraintGuard x ->
      hang
        (text "At" <+> pp (srcRange x) <.> colon)
        2
        (vcat [ text "Local declaration" <+> backticks (pp (thing x)) <+>
                text "may not use constraint guards."
              , text "Constraint guards may only appear at the top-level of" <+>
                text "a module."
              ])

expandPropGuards :: ModuleG mname PName -> ExpandPropGuardsM (ModuleG mname PName)
expandPropGuards m =
  do def <- expandModuleDef (mDef m)
     pure m { mDef = def }

expandModuleDef :: ModuleDefinition PName -> ExpandPropGuardsM (ModuleDefinition PName)
expandModuleDef m =
  case m of
    NormalModule ds    -> NormalModule . concat <$> mapM expandTopDecl ds
    FunctorInstance {} -> pure m
    InterfaceModule {} -> pure m

expandTopDecl :: TopDecl PName -> ExpandPropGuardsM [TopDecl PName]
expandTopDecl topDecl =
  case topDecl of
    Decl topLevelDecl ->
      do ds <- expandDecl (tlValue topLevelDecl)
         pure [ Decl topLevelDecl { tlValue = d } | d <- ds ]

    DModule tl | NestedModule m <- tlValue tl ->
      do m1 <- expandPropGuards m
         pure [DModule tl { tlValue = NestedModule m1 }]

    _ -> pure [topDecl]

expandDecl :: Decl PName -> ExpandPropGuardsM [Decl PName]
expandDecl decl =
  case decl of
    DBind bind ->
      do bs <- expandBind bind
         pure (map DBind bs)
    DLocated d rng ->
      do ds <- expandDecl d
         pure $ map (`DLocated` rng) ds
    _ ->
      pure [decl]

expandBind :: Bind PName -> ExpandPropGuardsM [Bind PName]
expandBind bind =
  case thing (bDef bind) of
    DImpl (DPropGuards guards) -> expand (DImpl . DPropGuards) guards
    DForeign (Just (DPropGuards guards)) -> expand (DForeign . Just . DPropGuards) guards
    _ ->
      do checkNestedGuardsInBind bind
         pure [bind]

  where
  expand def guards = do
    Forall params props t rng <-
      case bSignature bind of
        Just schema -> pure schema
        Nothing -> Left . NoSignature $ bName bind
    let goGuard ::
          PropGuardCase PName ->
          ExpandPropGuardsM (PropGuardCase PName, Bind PName)
        goGuard (PropGuardCase props' e) = do
          checkNestedGuardsInExpr e
          bName' <- newName (bName bind) (thing <$> props')
          -- call to generated function
          tParams <- case bSignature bind of
            Just (Forall tps _ _ _) -> pure tps
            Nothing -> Left $ NoSignature (bName bind)
          typeInsts <-
            (\(TParam n _ _) -> Right . PosInst $ TUser n [])
              `traverse` tParams
          let e' = foldl EApp (EAppT (EVar $ thing bName') typeInsts) (patternToExpr <$> bParams bind)
          pure
            ( PropGuardCase props' e',
              bind
                { bName = bName',
                  -- include guarded props in signature
                  bSignature = Just (Forall params
                                        (props <> map thing props')
                                        t rng),
                  -- keeps same location at original bind
                  -- i.e. "on top of" original bind
                  bDef = (bDef bind) {thing = exprDef e}
                }
            )
    (guards', binds') <- unzip <$> mapM goGuard guards
    pure $
      bind {bDef = def guards' <$ bDef bind} :
      binds'

-- | Check for nested uses of constraint guards in an expression (e.g.,
-- in a local definition bound with @where@).
checkNestedGuardsInExpr :: Expr PName -> ExpandPropGuardsM ()
checkNestedGuardsInExpr expr =
  case expr of
    EVar {} ->
      pure ()
    ELit {} ->
      pure ()
    EGenerate e ->
      checkNestedGuardsInExpr e
    ETuple es ->
      traverse_ checkNestedGuardsInExpr es
    ERecord r ->
      traverse_ (checkNestedGuardsInExpr . snd) r
    ESel e _ ->
      checkNestedGuardsInExpr e
    EUpd mbE upds ->
      do traverse_ checkNestedGuardsInExpr mbE
         traverse_ checkNestedGuardsInUpdField upds
    EList es ->
      traverse_ checkNestedGuardsInExpr es
    EFromTo {} ->
      pure ()
    EFromToBy {} ->
      pure ()
    EFromToDownBy {} ->
      pure ()
    EFromToLessThan {} ->
      pure ()
    EInfFrom _ mbE ->
      traverse_ checkNestedGuardsInExpr mbE
    EComp e mss ->
      do checkNestedGuardsInExpr e
         traverse_ (traverse_ checkNestedGuardsInMatch) mss
    EApp e1 e2 ->
      do checkNestedGuardsInExpr e1
         checkNestedGuardsInExpr e2
    EAppT e _ ->
      checkNestedGuardsInExpr e
    EIf e1 e2 e3 ->
      do checkNestedGuardsInExpr e1
         checkNestedGuardsInExpr e2
         checkNestedGuardsInExpr e3
    EWhere e ds ->
      do checkNestedGuardsInExpr e
         traverse_ checkNestedGuardsInDecl ds
    ETyped e _ ->
      checkNestedGuardsInExpr e
    ETypeVal _ ->
      pure ()
    EFun _ _ e ->
      checkNestedGuardsInExpr e
    ELocated e _ ->
      checkNestedGuardsInExpr e
    ESplit e ->
      checkNestedGuardsInExpr e
    EParens e ->
      checkNestedGuardsInExpr e
    EInfix e1 _ _ e2 ->
      do checkNestedGuardsInExpr e1
         checkNestedGuardsInExpr e2
    EPrefix _ e ->
      checkNestedGuardsInExpr e
    ECase e ca ->
      do checkNestedGuardsInExpr e
         traverse_ checkNestedGuardsInCaseAlt ca
  where
    checkNestedGuardsInUpdField :: UpdField PName -> ExpandPropGuardsM ()
    checkNestedGuardsInUpdField (UpdField _ _ e) =
      checkNestedGuardsInExpr e

    checkNestedGuardsInMatch :: Match PName -> ExpandPropGuardsM ()
    checkNestedGuardsInMatch match =
      case match of
        Match _ e ->
          checkNestedGuardsInExpr e
        MatchLet b ->
          checkNestedGuardsInBind b

    checkNestedGuardsInCaseAlt :: CaseAlt PName -> ExpandPropGuardsM ()
    checkNestedGuardsInCaseAlt (CaseAlt _ e) =
      checkNestedGuardsInExpr e

-- | Check for nested uses of constraint guards in a local declaration.
checkNestedGuardsInDecl :: Decl PName -> ExpandPropGuardsM ()
checkNestedGuardsInDecl decl =
  case decl of
    DSignature {} -> pure ()
    DPatBind _ e  -> checkNestedGuardsInExpr e
    DBind b       -> checkNestedGuardsInBind b
    DRec bs       -> traverse_ checkNestedGuardsInBind bs
    DFixity {}    -> pure ()
    DPragma {}    -> pure ()
    DType {}      -> pure ()
    DProp {}      -> pure ()
    DLocated d _  -> checkNestedGuardsInDecl d

-- | Check for nested uses of constraint guards in a local bind.
checkNestedGuardsInBind :: Bind PName -> ExpandPropGuardsM ()
checkNestedGuardsInBind bind =
  case thing (bDef bind) of
    DPrim         -> pure ()
    DImpl bi      -> checkNestedGuardsInBindImpl bi
    DForeign mbBi -> traverse_ checkNestedGuardsInBindImpl mbBi
  where
    nestedConstraintGuards :: ExpandPropGuardsM ()
    nestedConstraintGuards = Left . NestedConstraintGuard $ bName bind

    checkNestedGuardsInBindImpl :: BindImpl PName -> ExpandPropGuardsM ()
    checkNestedGuardsInBindImpl bi =
      case bi of
        DPropGuards _ ->
          nestedConstraintGuards
        DExpr e ->
          checkNestedGuardsInExpr e

patternToExpr :: Pattern PName -> Expr PName
patternToExpr (PVar locName) = EVar (thing locName)
patternToExpr _ =
  panic "patternToExpr"
    ["Unimplemented: patternToExpr of anything other than PVar"]

newName :: Located PName -> [Prop PName] -> ExpandPropGuardsM (Located PName)
newName locName props =
  pure case thing locName of
    Qual modName ident ->
      let txt = identText ident
          txt' = pack $ show $ pp props
       in Qual modName (mkIdent $ txt <> txt') <$ locName
    UnQual ident ->
      let txt = identText ident
          txt' = pack $ show $ pp props
       in UnQual (mkIdent $ txt <> txt') <$ locName
    NewName _ _ ->
      panic "mkName"
        [ "During expanding prop guards"
        , "tried to make new name from NewName case of PName"
        ]