{-# LANGUAGE PatternSynonyms #-}

-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Helper functions for working with @Terms@ and @Syntax@
module Swarm.Language.Syntax.Util (
  mkOp,
  mkOp',
  unfoldApps,
  mkTuple,
  unTuple,

  -- * Erasure
  erase,
  eraseS,

  -- * Term traversal
  freeVarsS,
  freeVarsT,
  freeVarsV,
  mapFreeS,
  locVarToSyntax',
  asTree,
  measureAstSize,
) where

import Control.Lens (Traversal', para, universe, (%~), (^.), pattern Empty)
import Control.Monad (void)
import Data.Data (Data)
import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty qualified as NonEmpty
import Data.Set qualified as S
import Data.Tree
import Swarm.Language.Syntax.AST
import Swarm.Language.Syntax.Constants
import Swarm.Language.Syntax.Loc
import Swarm.Language.Syntax.Pattern
import Swarm.Language.Var (Var)

-- Setup for doctests

-- $setup
-- >>> import Control.Lens ((^.))
-- >>> import Swarm.Language.Syntax.Constants
-- >>> import Swarm.Language.Syntax.Loc
-- >>> import Swarm.Language.Syntax.Pattern
-- >>> import Swarm.Language.Syntax.AST

-- | Make an infix operation (e.g. @2 + 3@) a curried function
--   application (e.g. @((+) 2) 3@).
mkOp :: Const -> (SrcLoc, t) -> Syntax -> Syntax -> Syntax
mkOp :: forall t. Const -> (SrcLoc, t) -> Syntax -> Syntax -> Syntax
mkOp Const
c (SrcLoc
opLoc, t
_) s1 :: Syntax
s1@(Syntax SrcLoc
l1 Term
_) s2 :: Syntax
s2@(Syntax SrcLoc
l2 Term
_) = SrcLoc -> Term -> Syntax
Syntax SrcLoc
newLoc Term
newTerm
 where
  -- The new syntax spans all terms
  newLoc :: SrcLoc
newLoc = SrcLoc
l1 SrcLoc -> SrcLoc -> SrcLoc
forall a. Semigroup a => a -> a -> a
<> SrcLoc
opLoc SrcLoc -> SrcLoc -> SrcLoc
forall a. Semigroup a => a -> a -> a
<> SrcLoc
l2
  sop :: Syntax
sop = SrcLoc -> Term -> Syntax
Syntax SrcLoc
opLoc (Const -> Term
forall ty. Const -> Term' ty
TConst Const
c)
  newTerm :: Term
newTerm = Syntax -> Syntax -> Term
forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SApp (SrcLoc -> Term -> Syntax
Syntax (SrcLoc
l1 SrcLoc -> SrcLoc -> SrcLoc
forall a. Semigroup a => a -> a -> a
<> SrcLoc
opLoc) (Term -> Syntax) -> Term -> Syntax
forall a b. (a -> b) -> a -> b
$ Syntax -> Syntax -> Term
forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SApp Syntax
sop Syntax
s1) Syntax
s2

-- | Make an infix operation, discarding any location information
mkOp' :: Const -> Term -> Term -> Term
mkOp' :: Const -> Term -> Term -> Term
mkOp' Const
c Term
t1 = Term -> Term -> Term
TApp (Term -> Term -> Term
TApp (Const -> Term
forall ty. Const -> Term' ty
TConst Const
c) Term
t1)

-- | Turn function application chain into a list.
--
-- >>> syntaxWrap f = fmap (^. sTerm) . f . Syntax NoLoc
-- >>> syntaxWrap unfoldApps (mkOp' Mul (TInt 1) (TInt 2)) -- 1 * 2
-- TConst Mul :| [TInt 1,TInt 2]
unfoldApps :: Syntax' ty -> NonEmpty (Syntax' ty)
unfoldApps :: forall ty. Syntax' ty -> NonEmpty (Syntax' ty)
unfoldApps Syntax' ty
trm = NonEmpty (Syntax' ty) -> NonEmpty (Syntax' ty)
forall a. NonEmpty a -> NonEmpty a
NonEmpty.reverse (NonEmpty (Syntax' ty) -> NonEmpty (Syntax' ty))
-> ((Syntax' ty -> (Syntax' ty, Maybe (Syntax' ty)))
    -> NonEmpty (Syntax' ty))
-> (Syntax' ty -> (Syntax' ty, Maybe (Syntax' ty)))
-> NonEmpty (Syntax' ty)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Syntax' ty -> (Syntax' ty, Maybe (Syntax' ty)))
 -> Syntax' ty -> NonEmpty (Syntax' ty))
-> Syntax' ty
-> (Syntax' ty -> (Syntax' ty, Maybe (Syntax' ty)))
-> NonEmpty (Syntax' ty)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Syntax' ty -> (Syntax' ty, Maybe (Syntax' ty)))
-> Syntax' ty -> NonEmpty (Syntax' ty)
forall a b. (a -> (b, Maybe a)) -> a -> NonEmpty b
NonEmpty.unfoldr Syntax' ty
trm ((Syntax' ty -> (Syntax' ty, Maybe (Syntax' ty)))
 -> NonEmpty (Syntax' ty))
-> (Syntax' ty -> (Syntax' ty, Maybe (Syntax' ty)))
-> NonEmpty (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ \case
  Syntax' SrcLoc
_ (SApp Syntax' ty
s1 Syntax' ty
s2) Comments
_ ty
_ -> (Syntax' ty
s2, Syntax' ty -> Maybe (Syntax' ty)
forall a. a -> Maybe a
Just Syntax' ty
s1)
  Syntax' ty
s -> (Syntax' ty
s, Maybe (Syntax' ty)
forall a. Maybe a
Nothing)

-- | Create an appropriate `Term` out of a list of syntax nodes which
--   were enclosed with parentheses (and separated by commas).
mkTuple :: [Syntax] -> Term
-- () = TUnit
mkTuple :: [Syntax] -> Term
mkTuple [] = Term
forall ty. Term' ty
TUnit
-- (x) = x, but record the fact that it was explicitly parenthesized,
-- for better source location tracking
mkTuple [Syntax
x] = Syntax -> Term
forall ty. Syntax' ty -> Term' ty
SParens Syntax
x
-- (x,y) = SPair
mkTuple [Syntax
x, Syntax
y] = Syntax -> Syntax -> Term
forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SPair Syntax
x Syntax
y
-- (x,y,...) = recursively nested pairs.  Note that we do not assign
-- source spans to the nested tuples since they don't really come from
-- a specific place in the source.
mkTuple (Syntax
x : [Syntax]
r) = Syntax -> Syntax -> Term
forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SPair Syntax
x (SrcLoc -> Term -> Syntax
Syntax SrcLoc
NoLoc ([Syntax] -> Term
mkTuple [Syntax]
r))

-- | Decompose a nested tuple into a list of components.
unTuple :: Syntax' ty -> [Syntax' ty]
unTuple :: forall ty. Syntax' ty -> [Syntax' ty]
unTuple = \case
  Syntax' SrcLoc
_ (SPair Syntax' ty
s1 Syntax' ty
s2) Comments
_ ty
_ -> Syntax' ty
s1 Syntax' ty -> [Syntax' ty] -> [Syntax' ty]
forall a. a -> [a] -> [a]
: Syntax' ty -> [Syntax' ty]
forall ty. Syntax' ty -> [Syntax' ty]
unTuple Syntax' ty
s2
  Syntax' ty
s -> [Syntax' ty
s]

------------------------------------------------------------
-- Type erasure
------------------------------------------------------------

-- | Erase the type annotations from a 'Syntax' or 'Term' tree.
erase :: Functor t => t ty -> t ()
erase :: forall (t :: * -> *) ty. Functor t => t ty -> t ()
erase = t ty -> t ()
forall (t :: * -> *) ty. Functor t => t ty -> t ()
void

-- | Erase all annotations from a 'Syntax' node, turning it into a
--   bare 'Term'.
eraseS :: Syntax' ty -> Term
eraseS :: forall ty. Syntax' ty -> Term
eraseS (Syntax' SrcLoc
_ Term' ty
t Comments
_ ty
_) = Term' ty -> Term
forall (t :: * -> *) ty. Functor t => t ty -> t ()
erase Term' ty
t

------------------------------------------------------------
-- Free variable traversals
------------------------------------------------------------

-- | Traversal over those subterms of a term which represent free
--   variables.  The S suffix indicates that it is a `Traversal' over
--   the `Syntax` nodes (which contain type and source location info)
--   containing free variables inside a larger `Syntax` value.  Note
--   that if you want to get the list of all `Syntax` nodes
--   representing free variables, you can do so via @'toListOf'
--   'freeVarsS'@.
freeVarsS :: forall ty. Traversal' (Syntax' ty) (Syntax' ty)
freeVarsS :: forall ty (f :: * -> *).
Applicative f =>
(Syntax' ty -> f (Syntax' ty)) -> Syntax' ty -> f (Syntax' ty)
freeVarsS Syntax' ty -> f (Syntax' ty)
f = Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
forall a. Set a
S.empty
 where
  -- go :: Applicative f => Set Var -> Syntax' ty -> f (Syntax' ty)
  go :: Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound s :: Syntax' ty
s@(Syntax' SrcLoc
l Term' ty
t Comments
ty ty
cmts) = case Term' ty
t of
    Term' ty
TUnit -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TConst {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TDir {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TInt {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TAntiInt {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TText {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TAntiText {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TAntiSyn {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TBool {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TRobot {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TRef {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TRequire {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    TStock {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
    SRequirements Var
x Syntax' ty
s1 -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Var -> Syntax' ty -> Term' ty
forall ty. Var -> Syntax' ty -> Term' ty
SRequirements Var
x (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s1
    TVar Var
x
      | Var
x Var -> Set Var -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Var
bound -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
      | Bool
otherwise -> Syntax' ty -> f (Syntax' ty)
f Syntax' ty
s
    SLam LocVar
x Maybe Type
xty Syntax' ty
s1 -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ LocVar -> Maybe Type -> Syntax' ty -> Term' ty
forall ty. LocVar -> Maybe Type -> Syntax' ty -> Term' ty
SLam LocVar
x Maybe Type
xty (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert (LocVar -> Var
forall v. Located v -> v
lvVar LocVar
x) Set Var
bound) Syntax' ty
s1
    SApp Syntax' ty
s1 Syntax' ty
s2 -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Syntax' ty -> Syntax' ty -> Term' ty
forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SApp (Syntax' ty -> Syntax' ty -> Term' ty)
-> f (Syntax' ty) -> f (Syntax' ty -> Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s1 f (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s2
    SLet LetSyntax
ls Bool
r LocVar
x Maybe RawPolytype
xty Maybe Polytype
xpty Maybe Requirements
xreq Syntax' ty
s1 Syntax' ty
s2 ->
      let bound' :: Set Var
bound' = Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert (LocVar -> Var
forall v. Located v -> v
lvVar LocVar
x) Set Var
bound
       in f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ LetSyntax
-> Bool
-> LocVar
-> Maybe RawPolytype
-> Maybe Polytype
-> Maybe Requirements
-> Syntax' ty
-> Syntax' ty
-> Term' ty
forall ty.
LetSyntax
-> Bool
-> LocVar
-> Maybe RawPolytype
-> Maybe Polytype
-> Maybe Requirements
-> Syntax' ty
-> Syntax' ty
-> Term' ty
SLet LetSyntax
ls Bool
r LocVar
x Maybe RawPolytype
xty Maybe Polytype
xpty Maybe Requirements
xreq (Syntax' ty -> Syntax' ty -> Term' ty)
-> f (Syntax' ty) -> f (Syntax' ty -> Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound' Syntax' ty
s1 f (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound' Syntax' ty
s2
    STydef Located TDVar
x Polytype
xdef Maybe TydefInfo
tdInfo Syntax' ty
t1 -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Located TDVar
-> Polytype -> Maybe TydefInfo -> Syntax' ty -> Term' ty
forall ty.
Located TDVar
-> Polytype -> Maybe TydefInfo -> Syntax' ty -> Term' ty
STydef Located TDVar
x Polytype
xdef Maybe TydefInfo
tdInfo (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
t1
    SPair Syntax' ty
s1 Syntax' ty
s2 -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Syntax' ty -> Syntax' ty -> Term' ty
forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SPair (Syntax' ty -> Syntax' ty -> Term' ty)
-> f (Syntax' ty) -> f (Syntax' ty -> Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s1 f (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s2
    SBind Maybe LocVar
mx Maybe ty
mty Maybe Polytype
mpty Maybe Requirements
mreq Syntax' ty
s1 Syntax' ty
s2 -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Maybe LocVar
-> Maybe ty
-> Maybe Polytype
-> Maybe Requirements
-> Syntax' ty
-> Syntax' ty
-> Term' ty
forall ty.
Maybe LocVar
-> Maybe ty
-> Maybe Polytype
-> Maybe Requirements
-> Syntax' ty
-> Syntax' ty
-> Term' ty
SBind Maybe LocVar
mx Maybe ty
mty Maybe Polytype
mpty Maybe Requirements
mreq (Syntax' ty -> Syntax' ty -> Term' ty)
-> f (Syntax' ty) -> f (Syntax' ty -> Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s1 f (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set Var -> Syntax' ty -> f (Syntax' ty)
go ((Set Var -> Set Var)
-> (LocVar -> Set Var -> Set Var)
-> Maybe LocVar
-> Set Var
-> Set Var
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set Var -> Set Var
forall a. a -> a
id (Var -> Set Var -> Set Var
forall a. Ord a => a -> Set a -> Set a
S.insert (Var -> Set Var -> Set Var)
-> (LocVar -> Var) -> LocVar -> Set Var -> Set Var
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LocVar -> Var
forall v. Located v -> v
lvVar) Maybe LocVar
mx Set Var
bound) Syntax' ty
s2
    SDelay Syntax' ty
s1 -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Syntax' ty -> Term' ty
forall ty. Syntax' ty -> Term' ty
SDelay (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s1
    SRcd Map Var (Maybe (Syntax' ty))
m -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Map Var (Maybe (Syntax' ty)) -> Term' ty
forall ty. Map Var (Maybe (Syntax' ty)) -> Term' ty
SRcd (Map Var (Maybe (Syntax' ty)) -> Term' ty)
-> f (Map Var (Maybe (Syntax' ty))) -> f (Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Maybe (Syntax' ty) -> f (Maybe (Syntax' ty)))
-> Map Var (Maybe (Syntax' ty)) -> f (Map Var (Maybe (Syntax' ty)))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Var a -> f (Map Var b)
traverse ((Maybe (Syntax' ty) -> f (Maybe (Syntax' ty)))
 -> Map Var (Maybe (Syntax' ty))
 -> f (Map Var (Maybe (Syntax' ty))))
-> ((Syntax' ty -> f (Syntax' ty))
    -> Maybe (Syntax' ty) -> f (Maybe (Syntax' ty)))
-> (Syntax' ty -> f (Syntax' ty))
-> Map Var (Maybe (Syntax' ty))
-> f (Map Var (Maybe (Syntax' ty)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Syntax' ty -> f (Syntax' ty))
-> Maybe (Syntax' ty) -> f (Maybe (Syntax' ty))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse) (Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound) Map Var (Maybe (Syntax' ty))
m
    SProj Syntax' ty
s1 Var
x -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Syntax' ty -> Var -> Term' ty
forall ty. Syntax' ty -> Var -> Term' ty
SProj (Syntax' ty -> Var -> Term' ty)
-> f (Syntax' ty) -> f (Var -> Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s1 f (Var -> Term' ty) -> f Var -> f (Term' ty)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Var -> f Var
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Var
x
    SAnnotate Syntax' ty
s1 RawPolytype
pty -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Syntax' ty -> RawPolytype -> Term' ty
forall ty. Syntax' ty -> RawPolytype -> Term' ty
SAnnotate (Syntax' ty -> RawPolytype -> Term' ty)
-> f (Syntax' ty) -> f (RawPolytype -> Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s1 f (RawPolytype -> Term' ty) -> f RawPolytype -> f (Term' ty)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RawPolytype -> f RawPolytype
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure RawPolytype
pty
    SSuspend Syntax' ty
s1 -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Syntax' ty -> Term' ty
forall ty. Syntax' ty -> Term' ty
SSuspend (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s1
    SParens Syntax' ty
s1 -> f (Term' ty) -> f (Syntax' ty)
forall {f :: * -> *}.
Applicative f =>
f (Term' ty) -> f (Syntax' ty)
rewrap (f (Term' ty) -> f (Syntax' ty)) -> f (Term' ty) -> f (Syntax' ty)
forall a b. (a -> b) -> a -> b
$ Syntax' ty -> Term' ty
forall ty. Syntax' ty -> Term' ty
SParens (Syntax' ty -> Term' ty) -> f (Syntax' ty) -> f (Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set Var -> Syntax' ty -> f (Syntax' ty)
go Set Var
bound Syntax' ty
s1
    TType {} -> Syntax' ty -> f (Syntax' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Syntax' ty
s
   where
    rewrap :: f (Term' ty) -> f (Syntax' ty)
rewrap f (Term' ty)
s' = SrcLoc -> Term' ty -> Comments -> ty -> Syntax' ty
forall ty. SrcLoc -> Term' ty -> Comments -> ty -> Syntax' ty
Syntax' SrcLoc
l (Term' ty -> Comments -> ty -> Syntax' ty)
-> f (Term' ty) -> f (Comments -> ty -> Syntax' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Term' ty)
s' f (Comments -> ty -> Syntax' ty)
-> f Comments -> f (ty -> Syntax' ty)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Comments -> f Comments
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Comments
ty f (ty -> Syntax' ty) -> f ty -> f (Syntax' ty)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ty -> f ty
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ty
cmts

-- | Like 'freeVarsS', but traverse over the 'Term's containing free
--   variables.  More direct if you don't need to know the types or
--   source locations of the variables.  Note that if you want to get
--   the list of all `Term`s representing free variables, you can do
--   so via @'toListOf' 'freeVarsT'@.
freeVarsT :: forall ty. Traversal' (Syntax' ty) (Term' ty)
freeVarsT :: forall ty (f :: * -> *).
Applicative f =>
(Term' ty -> f (Term' ty)) -> Syntax' ty -> f (Syntax' ty)
freeVarsT = (Syntax' ty -> f (Syntax' ty)) -> Syntax' ty -> f (Syntax' ty)
forall ty (f :: * -> *).
Applicative f =>
(Syntax' ty -> f (Syntax' ty)) -> Syntax' ty -> f (Syntax' ty)
freeVarsS ((Syntax' ty -> f (Syntax' ty)) -> Syntax' ty -> f (Syntax' ty))
-> ((Term' ty -> f (Term' ty)) -> Syntax' ty -> f (Syntax' ty))
-> (Term' ty -> f (Term' ty))
-> Syntax' ty
-> f (Syntax' ty)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Term' ty -> f (Term' ty)) -> Syntax' ty -> f (Syntax' ty)
forall ty (f :: * -> *).
Functor f =>
(Term' ty -> f (Term' ty)) -> Syntax' ty -> f (Syntax' ty)
sTerm

-- | Traversal over the free variables of a term.  Like 'freeVarsS'
--   and 'freeVarsT', but traverse over the variable names themselves.
--   Note that if you want to get the set of all free variable names,
--   you can do so via @'Data.Set.Lens.setOf' 'freeVarsV'@.
freeVarsV :: Traversal' (Syntax' ty) Var
freeVarsV :: forall ty (f :: * -> *).
Applicative f =>
(Var -> f Var) -> Syntax' ty -> f (Syntax' ty)
freeVarsV = (Term' ty -> f (Term' ty)) -> Syntax' ty -> f (Syntax' ty)
forall ty (f :: * -> *).
Applicative f =>
(Term' ty -> f (Term' ty)) -> Syntax' ty -> f (Syntax' ty)
freeVarsT ((Term' ty -> f (Term' ty)) -> Syntax' ty -> f (Syntax' ty))
-> ((Var -> f Var) -> Term' ty -> f (Term' ty))
-> (Var -> f Var)
-> Syntax' ty
-> f (Syntax' ty)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\Var -> f Var
f -> \case TVar Var
x -> Var -> Term' ty
forall ty. Var -> Term' ty
TVar (Var -> Term' ty) -> f Var -> f (Term' ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Var -> f Var
f Var
x; Term' ty
t -> Term' ty -> f (Term' ty)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term' ty
t)

-- | Apply a function to all free occurrences of a particular
--   variable.
mapFreeS :: Var -> (Syntax' ty -> Syntax' ty) -> Syntax' ty -> Syntax' ty
mapFreeS :: forall ty.
Var -> (Syntax' ty -> Syntax' ty) -> Syntax' ty -> Syntax' ty
mapFreeS Var
x Syntax' ty -> Syntax' ty
f = (Syntax' ty -> Identity (Syntax' ty))
-> Syntax' ty -> Identity (Syntax' ty)
forall ty (f :: * -> *).
Applicative f =>
(Syntax' ty -> f (Syntax' ty)) -> Syntax' ty -> f (Syntax' ty)
freeVarsS ((Syntax' ty -> Identity (Syntax' ty))
 -> Syntax' ty -> Identity (Syntax' ty))
-> (Syntax' ty -> Syntax' ty) -> Syntax' ty -> Syntax' ty
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (\Syntax' ty
t -> case Syntax' ty
t Syntax' ty
-> Getting (Term' ty) (Syntax' ty) (Term' ty) -> Term' ty
forall s a. s -> Getting a s a -> a
^. Getting (Term' ty) (Syntax' ty) (Term' ty)
forall ty (f :: * -> *).
Functor f =>
(Term' ty -> f (Term' ty)) -> Syntax' ty -> f (Syntax' ty)
sTerm of TVar Var
y | Var
y Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
x -> Syntax' ty -> Syntax' ty
f Syntax' ty
t; Term' ty
_ -> Syntax' ty
t)

-- | Transform the AST into a Tree datatype.  Useful for
--   pretty-printing (e.g. via "Data.Tree.drawTree").
asTree :: Data a => Syntax' a -> Tree (Syntax' a)
asTree :: forall a. Data a => Syntax' a -> Tree (Syntax' a)
asTree = (Syntax' a -> [Tree (Syntax' a)] -> Tree (Syntax' a))
-> Syntax' a -> Tree (Syntax' a)
forall a r. Plated a => (a -> [r] -> r) -> a -> r
para Syntax' a -> [Tree (Syntax' a)] -> Tree (Syntax' a)
forall a. a -> [Tree a] -> Tree a
Node

-- | Each constructor is a assigned a value of 1, plus
--   any recursive syntax it entails.
measureAstSize :: Data a => Syntax' a -> Int
measureAstSize :: forall a. Data a => Syntax' a -> Int
measureAstSize = [Syntax' a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Syntax' a] -> Int)
-> (Syntax' a -> [Syntax' a]) -> Syntax' a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Syntax' a -> Bool) -> [Syntax' a] -> [Syntax' a]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Syntax' a -> Bool) -> Syntax' a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Syntax' a -> Bool
forall a. Syntax' a -> Bool
isNoop) ([Syntax' a] -> [Syntax' a])
-> (Syntax' a -> [Syntax' a]) -> Syntax' a -> [Syntax' a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Syntax' a -> [Syntax' a]
forall a. Plated a => a -> [a]
universe

-- | Don't count "noop" nodes towards the code size.  They are usually
--   inserted automatically, either in @{}@ or after a bare @def@.
isNoop :: Syntax' a -> Bool
isNoop :: forall a. Syntax' a -> Bool
isNoop = \case
  Syntax' SrcLoc
_ (TConst Const
Noop) Comments
_ a
_ -> Bool
True
  Syntax' a
_ -> Bool
False

locVarToSyntax' :: LocVar -> ty -> Syntax' ty
locVarToSyntax' :: forall ty. LocVar -> ty -> Syntax' ty
locVarToSyntax' (LV SrcLoc
s Var
v) = SrcLoc -> Term' ty -> Comments -> ty -> Syntax' ty
forall ty. SrcLoc -> Term' ty -> Comments -> ty -> Syntax' ty
Syntax' SrcLoc
s (Var -> Term' ty
forall ty. Var -> Term' ty
TVar Var
v) Comments
forall s. AsEmpty s => s
Empty