-- |
-- Module      :  Swarm.Language.Elaborate
-- Copyright   :  Brent Yorgey
-- Maintainer  :  byorgey@gmail.com
--
-- 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, given its top-level type.  This currently performs such
--   operations as rewriting @if@ expressions and recursive let
--   expressions to use laziness appropriately.  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 :: Term -> Term
elaborate :: Term -> Term
elaborate =
  -- Wrap all *free* variables in 'Force'.  Free variables must be
  -- referring to a previous definition, which are all wrapped in
  -- 'TDelay'.
  (Traversal' Term Term
fvT forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ Term -> Term -> Term
TApp (Const -> Term
TConst Const
Force))
    -- Now do additional rewriting on all subterms.
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Plated a => (a -> a) -> a -> a
transform Term -> Term
rewrite
 where
  -- For recursive let bindings, rewrite any occurrences of x to
  -- (force x).  When interpreting t1, we will put a binding (x |->
  -- delay t1) in the context.
  rewrite :: Term -> Term
rewrite (TLet Bool
True Var
x Maybe Polytype
ty Term
t1 Term
t2) = Bool -> Var -> Maybe Polytype -> Term -> Term -> Term
TLet Bool
True Var
x Maybe Polytype
ty (Var -> Term -> Term
wrapForce Var
x Term
t1) (Var -> Term -> Term
wrapForce Var
x Term
t2)
  -- Rewrite any recursive occurrences of x inside t1 to (force x).
  -- When a TDef is encountered at runtime its body will immediately
  -- be wrapped in a VDelay. However, to make this work we also need
  -- to wrap all free variables in any term with 'force' --- since
  -- any such variables must in fact refer to things previously
  -- bound by 'def'.
  rewrite (TDef Bool
True Var
x Maybe Polytype
ty Term
t1) = Bool -> Var -> Maybe Polytype -> Term -> Term
TDef Bool
True Var
x Maybe Polytype
ty (Var -> (Term -> Term) -> Term -> Term
mapFree1 Var
x (Term -> Term -> Term
TApp (Const -> Term
TConst Const
Force)) Term
t1)
  -- Rewrite @f $ x@ to @f x@.
  rewrite (TApp (TApp (TConst Const
AppF) Term
r) Term
l) = Term -> Term -> Term
TApp Term
r Term
l
  -- Leave any other subterms alone.
  rewrite Term
t = Term
t

wrapForce :: Var -> Term -> Term
wrapForce :: Var -> Term -> Term
wrapForce Var
x = Var -> (Term -> Term) -> Term -> Term
mapFree1 Var
x (Term -> Term -> Term
TApp (Const -> Term
TConst Const
Force))