{-|
Copyright   : (C) 2021, QBayLogic B.V.
License     : BSD2 (see the file LICENSE)
Maintainer  : QBayLogic B.V. <devops@qbaylogic.com>

Utility class to extract free variables from data which has variables.
-}

{-# LANGUAGE FlexibleInstances #-}

module Clash.Core.HasFreeVars
  ( HasFreeVars(..)
  ) where

import Control.Lens as Lens (foldMapOf)
import Data.Monoid (All(..), Any(..))

import Clash.Core.FreeVars
import Clash.Core.Term (Term)
import Clash.Core.Type (Type(..))
import Clash.Core.Var (Var)
import Clash.Core.VarEnv

class HasFreeVars a where
  {-# MINIMAL freeVarsOf #-}

  freeVarsOf :: a -> VarSet

  {-# INLINE isClosed #-}
  -- | Something is closed if it has no free variables.
  -- This function may be replaced with a more efficient implementation.
  isClosed :: a -> Bool
  isClosed = VarSet -> Bool
nullVarSet (VarSet -> Bool) -> (a -> VarSet) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf

  {-# INLINE elemFreeVars #-}
  -- | Check if a variable is free in the given value.
  -- This function may be replaced with a more efficient implementation.
  elemFreeVars :: Var a -> a -> Bool
  elemFreeVars Var a
v = Var a -> VarSet -> Bool
forall a. Var a -> VarSet -> Bool
elemVarSet Var a
v (VarSet -> Bool) -> (a -> VarSet) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf

  {-# INLINE notElemFreeVars #-}
  -- | Check if a variable is not free in the given value.
  -- This function may be replaced with a more efficient implementation.
  notElemFreeVars :: Var a -> a -> Bool
  notElemFreeVars Var a
x = Var a -> VarSet -> Bool
forall a. Var a -> VarSet -> Bool
notElemVarSet Var a
x (VarSet -> Bool) -> (a -> VarSet) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf

  {-# INLINE subsetFreeVars #-}
  -- | Check if all variables in a set are free in the given value.
  -- This function may be replaced with a more efficient implementation.
  subsetFreeVars :: VarSet -> a -> Bool
  subsetFreeVars VarSet
xs = VarSet -> VarSet -> Bool
subsetVarSet VarSet
xs (VarSet -> Bool) -> (a -> VarSet) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf

  {-# INLINE disjointFreeVars #-}
  -- | Check if no variables in a set are free in the given value.
  -- This function may be replaced with a more efficient implementation.
  disjointFreeVars :: VarSet -> a -> Bool
  disjointFreeVars VarSet
xs = VarSet -> VarSet -> Bool
disjointVarSet VarSet
xs (VarSet -> Bool) -> (a -> VarSet) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf

instance HasFreeVars Term where
  {-# INLINE freeVarsOf #-}
  freeVarsOf :: Term -> VarSet
freeVarsOf =
    Getting VarSet Term (Var Any)
-> (Var Any -> VarSet) -> Term -> VarSet
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting VarSet Term (Var Any)
forall a. Fold Term (Var a)
freeLocalVars Var Any -> VarSet
forall a. Var a -> VarSet
unitVarSet

  elemFreeVars :: Var Term -> Term -> Bool
elemFreeVars Var Term
v Term
e =
    Any -> Bool
getAny (Getting Any Term (Var Term) -> (Var Term -> Any) -> Term -> Any
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting Any Term (Var Term)
forall a. Fold Term (Var a)
freeLocalVars (Bool -> Any
Any (Bool -> Any) -> (Var Term -> Bool) -> Var Term -> Any
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Var Term -> Var Term -> Bool
forall a. Eq a => a -> a -> Bool
== Var Term
v)) Term
e)

  notElemFreeVars :: Var Term -> Term -> Bool
notElemFreeVars Var Term
v Term
e =
    All -> Bool
getAll (Getting All Term (Var Term) -> (Var Term -> All) -> Term -> All
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting All Term (Var Term)
forall a. Fold Term (Var a)
freeLocalVars (Bool -> All
All (Bool -> All) -> (Var Term -> Bool) -> Var Term -> All
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Var Term -> Var Term -> Bool
forall a. Eq a => a -> a -> Bool
/= Var Term
v)) Term
e)

  disjointFreeVars :: VarSet -> Term -> Bool
disjointFreeVars VarSet
vs Term
e =
    All -> Bool
getAll (Getting All Term (Var Any) -> (Var Any -> All) -> Term -> All
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting All Term (Var Any)
forall a. Fold Term (Var a)
freeLocalVars (Bool -> All
All (Bool -> All) -> (Var Any -> Bool) -> Var Any -> All
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Var Any -> VarSet -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` VarSet
vs)) Term
e)

instance HasFreeVars Type where
  {-# INLINE freeVarsOf #-}
  freeVarsOf :: Type -> VarSet
freeVarsOf =
    Getting VarSet Type (Var Type)
-> (Var Type -> VarSet) -> Type -> VarSet
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting VarSet Type (Var Type)
Fold Type (Var Type)
typeFreeVars Var Type -> VarSet
forall a. Var a -> VarSet
unitVarSet

  isClosed :: Type -> Bool
isClosed Type
ty =
    case Type
ty of
      VarTy{} -> Bool
False
      ForAllTy{} -> All -> Bool
getAll (Getting All Type (Var Type) -> (Var Type -> All) -> Type -> All
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting All Type (Var Type)
Fold Type (Var Type)
typeFreeVars (All -> Var Type -> All
forall a b. a -> b -> a
const (Bool -> All
All Bool
False)) Type
ty)
      AppTy Type
l Type
r -> Type -> Bool
forall a. HasFreeVars a => a -> Bool
isClosed Type
l Bool -> Bool -> Bool
&& Type -> Bool
forall a. HasFreeVars a => a -> Bool
isClosed Type
r
      Type
_ -> Bool
True

  elemFreeVars :: Var Type -> Type -> Bool
elemFreeVars Var Type
v Type
ty =
    Any -> Bool
getAny (Getting Any Type (Var Type) -> (Var Type -> Any) -> Type -> Any
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting Any Type (Var Type)
Fold Type (Var Type)
typeFreeVars (Bool -> Any
Any (Bool -> Any) -> (Var Type -> Bool) -> Var Type -> Any
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Var Type -> Var Type -> Bool
forall a. Eq a => a -> a -> Bool
== Var Type
v)) Type
ty)

  notElemFreeVars :: Var Type -> Type -> Bool
notElemFreeVars Var Type
v Type
ty =
    All -> Bool
getAll (Getting All Type (Var Type) -> (Var Type -> All) -> Type -> All
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting All Type (Var Type)
Fold Type (Var Type)
typeFreeVars (Bool -> All
All (Bool -> All) -> (Var Type -> Bool) -> Var Type -> All
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Var Type -> Var Type -> Bool
forall a. Eq a => a -> a -> Bool
/= Var Type
v)) Type
ty)

instance (Foldable f, HasFreeVars a) => HasFreeVars (f a) where
  {-# INLINE freeVarsOf #-}
  freeVarsOf :: f a -> VarSet
freeVarsOf = (a -> VarSet) -> f a -> VarSet
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> VarSet
forall a. HasFreeVars a => a -> VarSet
freeVarsOf