{-# LANGUAGE OverloadedStrings #-}

-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Term elaboration which happens after type checking.
module Swarm.Language.Elaborate where

import Control.Lens (transform, (^.))
import Swarm.Language.Syntax

-- | Perform some elaboration / rewriting on a fully type-annotated
--   term.  This currently performs such operations as rewriting @if@
--   expressions and recursive let expressions to use laziness
--   appropriately.  It also inserts types and requirements on bound
--   variables so they will be available at runtime.
--
--   In theory it could also perform rewriting for overloaded
--   constants depending on the actual type they are used at, but
--   currently that sort of thing tends to make type inference fall
--   over.
elaborate :: TSyntax -> TSyntax
elaborate :: TSyntax -> TSyntax
elaborate = (TSyntax -> TSyntax) -> TSyntax -> TSyntax
forall a. Plated a => (a -> a) -> a -> a
transform TSyntax -> TSyntax
rewrite
 where
  rewrite :: TSyntax -> TSyntax
  rewrite :: TSyntax -> TSyntax
rewrite (Syntax' SrcLoc
l Term' (Poly 'Quantified Type)
t Comments
cs Poly 'Quantified Type
ty) = SrcLoc
-> Term' (Poly 'Quantified Type)
-> Comments
-> Poly 'Quantified Type
-> TSyntax
forall ty. SrcLoc -> Term' ty -> Comments -> ty -> Syntax' ty
Syntax' SrcLoc
l (Term' (Poly 'Quantified Type) -> Term' (Poly 'Quantified Type)
rewriteTerm Term' (Poly 'Quantified Type)
t) Comments
cs Poly 'Quantified Type
ty

  rewriteTerm :: TTerm -> TTerm
  rewriteTerm :: Term' (Poly 'Quantified Type) -> Term' (Poly 'Quantified Type)
rewriteTerm = \case
    -- Here we take inferred types for variables bound by def or
    -- bind (but NOT let) and stuff them into the term itself, so that
    -- we will still have access to them at runtime, after type
    -- annotations on the AST are erased.  We need them at runtime so
    -- we can keep track of the types of variables in scope, for use
    -- in typechecking additional terms entered at the REPL.  The
    -- reason we do not do this for 'let' is so that 'let' introduces
    -- truly local bindings which will not be available for use in
    -- later REPL terms.
    --
    -- We assume requirements for these variables have already been
    -- filled in during typechecking.  The reason we need to wait
    -- until now to fill in the types is that we have to wait until
    -- solving, substitution, and generalization are complete.
    --
    -- Eventually, once requirements analysis is part of typechecking,
    -- we'll infer them both at typechecking time then fill them in
    -- during elaboration here.
    SLet LetSyntax
ls Bool
r LocVar
x Maybe RawPolytype
mty Maybe (Poly 'Quantified Type)
_ Maybe Requirements
mreq TSyntax
t1 TSyntax
t2 ->
      let mpty :: Maybe (Poly 'Quantified Type)
mpty = case LetSyntax
ls of
            LetSyntax
LSDef -> Poly 'Quantified Type -> Maybe (Poly 'Quantified Type)
forall a. a -> Maybe a
Just (TSyntax
t1 TSyntax
-> Getting (Poly 'Quantified Type) TSyntax (Poly 'Quantified Type)
-> Poly 'Quantified Type
forall s a. s -> Getting a s a -> a
^. Getting (Poly 'Quantified Type) TSyntax (Poly 'Quantified Type)
forall ty (f :: * -> *).
Functor f =>
(ty -> f ty) -> Syntax' ty -> f (Syntax' ty)
sType)
            LetSyntax
LSLet -> Maybe (Poly 'Quantified Type)
forall a. Maybe a
Nothing
       in LetSyntax
-> Bool
-> LocVar
-> Maybe RawPolytype
-> Maybe (Poly 'Quantified Type)
-> Maybe Requirements
-> TSyntax
-> TSyntax
-> Term' (Poly 'Quantified Type)
forall ty.
LetSyntax
-> Bool
-> LocVar
-> Maybe RawPolytype
-> Maybe (Poly 'Quantified Type)
-> Maybe Requirements
-> Syntax' ty
-> Syntax' ty
-> Term' ty
SLet LetSyntax
ls Bool
r LocVar
x Maybe RawPolytype
mty Maybe (Poly 'Quantified Type)
mpty Maybe Requirements
mreq TSyntax
t1 TSyntax
t2
    SBind Maybe LocVar
x (Just Poly 'Quantified Type
ty) Maybe (Poly 'Quantified Type)
_ Maybe Requirements
mreq TSyntax
c1 TSyntax
c2 -> Maybe LocVar
-> Maybe (Poly 'Quantified Type)
-> Maybe (Poly 'Quantified Type)
-> Maybe Requirements
-> TSyntax
-> TSyntax
-> Term' (Poly 'Quantified Type)
forall ty.
Maybe LocVar
-> Maybe ty
-> Maybe (Poly 'Quantified Type)
-> Maybe Requirements
-> Syntax' ty
-> Syntax' ty
-> Term' ty
SBind Maybe LocVar
x Maybe (Poly 'Quantified Type)
forall a. Maybe a
Nothing (Poly 'Quantified Type -> Maybe (Poly 'Quantified Type)
forall a. a -> Maybe a
Just Poly 'Quantified Type
ty) Maybe Requirements
mreq TSyntax
c1 TSyntax
c2
    -- Rewrite @f $ x@ to @f x@.
    SApp (Syntax' SrcLoc
_ (SApp (Syntax' SrcLoc
_ (TConst Const
AppF) Comments
_ Poly 'Quantified Type
_) TSyntax
l) Comments
_ Poly 'Quantified Type
_) TSyntax
r -> TSyntax -> TSyntax -> Term' (Poly 'Quantified Type)
forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SApp TSyntax
l TSyntax
r
    -- Leave any other subterms alone.
    Term' (Poly 'Quantified Type)
t -> Term' (Poly 'Quantified Type)
t

-- | Insert a special 'suspend' primitive at the very end of an erased
--   term which must have a command type.
insertSuspend :: Term -> Term
insertSuspend :: Term -> Term
insertSuspend Term
t = case Term
t of
  -- Primitive things which have type Cmd Unit: p => (p ; suspend ())
  TRequire {} -> Term
thenSuspend
  TStock {} -> Term
thenSuspend
  TRequirements {} -> Term
thenSuspend
  -- Recurse through def, tydef, bind, and annotate.
  TLet LetSyntax
ls Bool
r Var
x Maybe RawPolytype
mty Maybe (Poly 'Quantified Type)
mpty Maybe Requirements
mreq Term
t1 Term
t2 -> LetSyntax
-> Bool
-> Var
-> Maybe RawPolytype
-> Maybe (Poly 'Quantified Type)
-> Maybe Requirements
-> Term
-> Term
-> Term
TLet LetSyntax
ls Bool
r Var
x Maybe RawPolytype
mty Maybe (Poly 'Quantified Type)
mpty Maybe Requirements
mreq Term
t1 (Term -> Term
insertSuspend Term
t2)
  TTydef TDVar
x Poly 'Quantified Type
pty Maybe TydefInfo
mtd Term
t1 -> TDVar -> Poly 'Quantified Type -> Maybe TydefInfo -> Term -> Term
TTydef TDVar
x Poly 'Quantified Type
pty Maybe TydefInfo
mtd (Term -> Term
insertSuspend Term
t1)
  TBind Maybe Var
mx Maybe (Poly 'Quantified Type)
mty Maybe Requirements
mreq Term
c1 Term
c2 -> Maybe Var
-> Maybe (Poly 'Quantified Type)
-> Maybe Requirements
-> Term
-> Term
-> Term
TBind Maybe Var
mx Maybe (Poly 'Quantified Type)
mty Maybe Requirements
mreq Term
c1 (Term -> Term
insertSuspend Term
c2)
  TAnnotate Term
t1 RawPolytype
ty -> Term -> RawPolytype -> Term
TAnnotate (Term -> Term
insertSuspend Term
t1) RawPolytype
ty
  -- Replace pure or noop with suspend
  TApp (TConst Const
Pure) Term
t1 -> Term -> Term
TSuspend Term
t1
  TConst Const
Noop -> Term -> Term
TSuspend Term
forall ty. Term' ty
TUnit
  -- Anything else: p => (__res__ <- p; suspend __res__)
  --
  -- Note that since we don't put type + reqs annotations on
  -- __res__, it won't get added to the type environment and hence
  -- won't be in scope for use after the suspend.  But we pick a
  -- weird unlikely-to-be-used name just to be safe.
  Term
t' -> Maybe Var
-> Maybe (Poly 'Quantified Type)
-> Maybe Requirements
-> Term
-> Term
-> Term
TBind (Var -> Maybe Var
forall a. a -> Maybe a
Just Var
"__res__") Maybe (Poly 'Quantified Type)
forall a. Maybe a
Nothing Maybe Requirements
forall a. Maybe a
Nothing Term
t' (Term -> Term
TSuspend (Var -> Term
forall ty. Var -> Term' ty
TVar Var
"__res__"))
 where
  thenSuspend :: Term
thenSuspend = Maybe Var
-> Maybe (Poly 'Quantified Type)
-> Maybe Requirements
-> Term
-> Term
-> Term
TBind Maybe Var
forall a. Maybe a
Nothing Maybe (Poly 'Quantified Type)
forall a. Maybe a
Nothing Maybe Requirements
forall a. Maybe a
Nothing Term
t (Term -> Term
TSuspend Term
forall ty. Term' ty
TUnit)