Classes.hs

-- |
-- Module      : Rebound.Classes
-- Description : Type class definitions
--
-- Main typeclasses used by the library.

{-# LANGUAGE DefaultSignatures #-}
module Rebound.Classes where

import Rebound.Lib
import Data.LocalName
import Data.Scoped.List(List, pattern Nil, pattern (:<))
import Data.Scoped.List qualified as List

import Data.Foldable
import Data.Vec qualified as Vec
import Data.Fin qualified as Fin
import Data.Set (Set)
import Data.Set qualified as Set

import GHC.Generics (Generic1(..))

----------------------------------------------------------
-- Indices/variables shifting
----------------------------------------------------------

-- | Bring a scoped type into a new, bigger, scope, through variable shifting.
-- 
-- This class is used for types which are scoped, yet do not allow
-- substitution in general. Typical examples are data-structures which
-- associate metadata to variables.
-- See 'Rebound.Refinement.Refinement' for an example. 
class Shiftable t where
  shift :: SNat k -> t n -> t (k + n)
  -- a good default implementation of this is `shiftFromApply`. But the 
  -- `Subst` class is not yet in scope.  
  
----------------------------------------------------------
-- Free variables
----------------------------------------------------------

-- | Computes the set of free variables in a term.
class FV (t :: Nat -> Type) where
  -- | Does a particular variable appear free?
  appearsFree :: Fin n -> t n -> Bool
  default appearsFree :: (Generic1 t, GFV (Rep1 t)) => Fin n -> t n -> Bool
  appearsFree x e = gappearsFree x (from1 e)
  {-# INLINE appearsFree #-}

  -- | Calculate all of the free variables in a term.
  freeVars :: t n -> Set (Fin n)
  default freeVars :: (Generic1 t, GFV (Rep1 t)) => t n -> Set (Fin n)
  freeVars e = gfreeVars (from1 e)
  {-# INLINE freeVars #-}

-- | Generic programming support for 'FV'.
class GFV (t :: Nat -> Type) where
  gappearsFree :: Fin n -> t n -> Bool
  gfreeVars :: t n -> Set (Fin n)

----------------------------------------------------------
-- * Strengthening
----------------------------------------------------------

-- Strengthening cannot be implemented through substitution because it
-- must fail if the term uses invalid variables. Therefore, we make a
-- class of scoped types that can be strengthened.

-- | Eliminates the most recently bound variable from the term (if unused).
strengthen :: forall n t. (Strengthen t, SNatI n) => t (S n) -> Maybe (t n)
strengthen = strengthenRec s0 s1 (snat :: SNat n)

-- | Eliminates the @n@ most recently bound variables from the term (if unused).
strengthenN :: forall m n t. (Strengthen t, SNatI n) => SNat m -> t (m + n) -> Maybe (t n)
strengthenN m = strengthenRec s0 m (snat :: SNat n)

-- | Bring scoped terms into a smaller scope, if possible.
--
-- Strengthening is only possible if the term only refers to variable which
-- are in the smaller scope.
class Strengthen t where
  -- generalize strengthening -- remove m variables from the middle of the scope
  strengthenRec :: SNat k -> SNat m -> SNat n -> t (k + (m + n)) -> Maybe (t (k + n))
  default strengthenRec :: (Generic1 t, GStrengthen (Rep1 t)) =>
     SNat k -> SNat m -> SNat n -> t (k + (m + n)) -> Maybe (t (k + n))
  strengthenRec k m n t = to1 <$> gstrengthenRec k m n (from1 t)

  -- Remove a single variable from the middle of the scope
  strengthenOneRec :: forall k n. SNat k -> SNat n -> t (k + S n) -> Maybe (t (k + n))
  strengthenOneRec k = strengthenRec k s1

-- | Generic programming support for 'Strengthen'.
class GStrengthen (t :: Nat -> Type) where
  gstrengthenRec :: SNat k -> SNat m -> SNat n -> t (k + (m + n)) -> Maybe (t (k + n))

----------------------------------------------------------
-- FV and Strengthen instances for Data.Scoped.List
---------------------------------------------------------

instance FV t => FV (List t) where
  appearsFree :: Fin n -> List t n -> Bool
  appearsFree x = List.any (appearsFree x)

  freeVars :: List t n -> Set (Fin n)
  freeVars = List.foldr (\x s -> freeVars x `Set.union` s) Set.empty

instance Strengthen t => Strengthen (List t) where
  strengthenRec :: SNat k -> SNat m -> SNat n -> List t (k + (m + n)) -> Maybe (List t (k + n))
  strengthenRec k m n Nil = Just Nil
  strengthenRec k m n (x :< xs) = (:<) <$> strengthenRec k m n x <*> strengthenRec k m n xs

----------------------------------------------------------
-- FV and Strengthen instances for Fin
---------------------------------------------------------

instance FV Fin where
  appearsFree = (==)
  freeVars = Set.singleton

instance Strengthen Fin where
  strengthenRec :: SNat k -> SNat m -> SNat n-> Fin (k + (m + n)) -> Maybe (Fin (k + n))
  strengthenRec = Fin.strengthenRecFin

-- | Update a set of free variables to a new scope through strengthening
rescope :: forall n k. SNat k -> Set (Fin (k + n)) -> Set (Fin n)
rescope k = foldMap g where
   g :: Fin (k + n) -> Set (Fin n)
   g x = maybe
     Set.empty Set.singleton
     (Fin.strengthenRecFin s0 k (undefined :: SNat n) x)

----------------------------------------------------------
-- Type classes for patterns
----------------------------------------------------------

-- | Calculate the number of binding variables in the pattern
-- This number does not need to be an explicit parameter of the type
-- so that we have flexibility about what types we can use as
-- patterns.
class Sized (t :: Type) where
  -- Retrieve size from the type (number of variables bound by the pattern)
  type Size t :: Nat
  -- Access size as a term
  size :: t -> SNat (Size t)

-- | Pairs of types that can be compared with each other as patterns
class PatEq (t1 :: Type) (t2 :: Type) where
    patEq :: t1 -> t2 -> Maybe (Size t1 :~: Size t2)

-- | Class of patterns that are indexed by a natural number
-- where the size is that index directly
class (Sized (t p), Size (t p) ~ p) => SizeIndex t p


---------------------------------------------------------
-- Pattern Class Instances for Prelude and Lib Types
---------------------------------------------------------

-- ** LocalNames

instance Sized LocalName where
  type Size LocalName = N1
  size _ = s1

instance PatEq LocalName LocalName where
  patEq p1 p2 = Just Refl

-- ** SNats
instance Sized (SNat n) where
  type Size (SNat n) = n
  size n = n

instance PatEq (SNat n1) (SNat n2) where
  patEq = testEquality


-- ** Vectors

instance Sized (Vec n a) where
  type Size (Vec n a) = n
  size = Vec.vlength

instance Eq a => PatEq (Vec n1 a) (Vec n2 a) where
  patEq VNil VNil = Just Refl
  patEq (x ::: xs) (y ::: ys) | x == y,
    Just Refl <- patEq xs ys
    = Just Refl
  patEq _ _ = Nothing

-- ** Unit (trivial)

instance Sized () where { type Size () = N0 ;  size _ = SZ }

instance PatEq () () where patEq _ _ = Just Refl

-- ** Pairs

instance (Sized a, Sized b) => Sized (a,b) where
   type Size (a,b) = Size a + Size b
   size (x,y) = sPlus (size x) (size y)

instance (PatEq a1 a2, PatEq b1 b2) => PatEq (a1, b1) (a2, b2) where
   patEq (x1,y1) (x2,y2)
     | Just Refl <- patEq x1 x2
     , Just Refl <- patEq y1 y2
     = Just Refl
   patEq _ _ = Nothing

------------------------------------------