{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
-- Arbitrary NonEmpty
{-# OPTIONS_GHC -Wno-orphans #-}

module Constrained.Core (
  Var (..),
  eqVar,
  Rename (rename),
  freshVar,
  freshen,
  Value (..),
  unValue,
  NonEmpty ((:|)),
  Evidence (..),
  unionWithMaybe,
) where

import Constrained.List (
  List (..),
  mapList,
 )
import Constrained.PrettyUtils
import Control.Applicative
import Data.Function
import Data.List.NonEmpty (NonEmpty ((:|)))
import Data.List.NonEmpty qualified as NE
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Typeable
import Test.QuickCheck (Arbitrary (..), NonEmptyList (NonEmpty))

-- Variables --------------------------------------------------------------

data Var a = Var {forall a. Var a -> Int
nameOf :: Int, forall a. Var a -> String
nameHint :: String}

instance Ord (Var a) where
  compare :: Var a -> Var a -> Ordering
compare = Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> (Var a -> Int) -> Var a -> Var a -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` Var a -> Int
forall a. Var a -> Int
nameOf

instance Eq (Var a) where
  == :: Var a -> Var a -> Bool
(==) = Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Int -> Int -> Bool) -> (Var a -> Int) -> Var a -> Var a -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` Var a -> Int
forall a. Var a -> Int
nameOf

instance Show (Var a) where
  show :: Var a -> String
show Var a
v = Var a -> String
forall a. Var a -> String
nameHint Var a
v String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Var a -> Int
forall a. Var a -> Int
nameOf Var a
v)

eqVar :: forall a a'. (Typeable a, Typeable a') => Var a -> Var a' -> Maybe (a :~: a')
eqVar :: forall a a'.
(Typeable a, Typeable a') =>
Var a -> Var a' -> Maybe (a :~: a')
eqVar Var a
v Var a'
v' | Var a -> Int
forall a. Var a -> Int
nameOf Var a
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Var a' -> Int
forall a. Var a -> Int
nameOf Var a'
v' = forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall a b. (Typeable a, Typeable b) => Maybe (a :~: b)
eqT @a @a'
eqVar Var a
_ Var a'
_ = Maybe (a :~: a')
forall a. Maybe a
Nothing

-- Variable renaming ------------------------------------------------------

class Rename a where
  rename :: Typeable x => Var x -> Var x -> a -> a

instance Typeable a => Rename (Var a) where
  rename :: forall x. Typeable x => Var x -> Var x -> Var a -> Var a
rename Var x
v Var x
v' Var a
v''
    | Just x :~: a
Refl <- Var x -> Var a -> Maybe (x :~: a)
forall a a'.
(Typeable a, Typeable a') =>
Var a -> Var a' -> Maybe (a :~: a')
eqVar Var x
v Var a
v'' = Var a
Var x
v'
    | Bool
otherwise = Var a
v''

instance Rename () where
  rename :: forall x. Typeable x => Var x -> Var x -> () -> ()
rename Var x
_ Var x
_ ()
_ = ()

instance (Rename a, Rename b) => Rename (a, b) where
  rename :: forall x. Typeable x => Var x -> Var x -> (a, b) -> (a, b)
rename Var x
x Var x
x' (a
a, b
b) = (Var x -> Var x -> a -> a
forall x. Typeable x => Var x -> Var x -> a -> a
forall a x. (Rename a, Typeable x) => Var x -> Var x -> a -> a
rename Var x
x Var x
x' a
a, Var x -> Var x -> b -> b
forall x. Typeable x => Var x -> Var x -> b -> b
forall a x. (Rename a, Typeable x) => Var x -> Var x -> a -> a
rename Var x
x Var x
x' b
b)

instance {-# OVERLAPPABLE #-} (Functor t, Rename a) => Rename (t a) where
  rename :: forall x. Typeable x => Var x -> Var x -> t a -> t a
rename Var x
v Var x
v'
    | Var x
v Var x -> Var x -> Bool
forall a. Eq a => a -> a -> Bool
== Var x
v' = t a -> t a
forall a. a -> a
id
    | Bool
otherwise = (a -> a) -> t a -> t a
forall a b. (a -> b) -> t a -> t b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Var x -> Var x -> a -> a
forall x. Typeable x => Var x -> Var x -> a -> a
forall a x. (Rename a, Typeable x) => Var x -> Var x -> a -> a
rename Var x
v Var x
v')

instance (Ord a, Rename a) => Rename (Set a) where
  rename :: forall x. Typeable x => Var x -> Var x -> Set a -> Set a
rename Var x
v Var x
v'
    | Var x
v Var x -> Var x -> Bool
forall a. Eq a => a -> a -> Bool
== Var x
v' = Set a -> Set a
forall a. a -> a
id
    | Bool
otherwise = (a -> a) -> Set a -> Set a
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map (Var x -> Var x -> a -> a
forall x. Typeable x => Var x -> Var x -> a -> a
forall a x. (Rename a, Typeable x) => Var x -> Var x -> a -> a
rename Var x
v Var x
v')

instance (forall a. Rename (f a)) => Rename (List f as) where
  rename :: forall x. Typeable x => Var x -> Var x -> List f as -> List f as
rename Var x
v Var x
v' = (forall a. f a -> f a) -> List f as -> List f as
forall {k} (f :: k -> *) (g :: k -> *) (as :: [k]).
(forall (a :: k). f a -> g a) -> List f as -> List g as
mapList (Var x -> Var x -> f a -> f a
forall x. Typeable x => Var x -> Var x -> f a -> f a
forall a x. (Rename a, Typeable x) => Var x -> Var x -> a -> a
rename Var x
v Var x
v')

freshVar :: Var a -> Set Int -> Var a
freshVar :: forall a. Var a -> Set Int -> Var a
freshVar (Var Int
n String
nh) Set Int
ns
  | Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Int
n Set Int
ns = Int -> String -> Var a
forall a. Int -> String -> Var a
Var (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Set Int -> Int
forall a. Set a -> a
Set.findMax Set Int
ns) String
nh
  | Bool
otherwise = Int -> String -> Var a
forall a. Int -> String -> Var a
Var Int
n String
nh

freshen :: (Typeable a, Rename t) => Var a -> t -> Set Int -> (Var a, t)
freshen :: forall a t.
(Typeable a, Rename t) =>
Var a -> t -> Set Int -> (Var a, t)
freshen Var a
v t
t Set Int
nms
  | Var a -> Int
forall a. Var a -> Int
nameOf Var a
v Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Int
nms = let v' :: Var a
v' = Var a -> Set Int -> Var a
forall a. Var a -> Set Int -> Var a
freshVar Var a
v Set Int
nms in (Var a
v', Var a -> Var a -> t -> t
forall x. Typeable x => Var x -> Var x -> t -> t
forall a x. (Rename a, Typeable x) => Var x -> Var x -> a -> a
rename Var a
v Var a
v' t
t)
  | Bool
otherwise = (Var a
v, t
t)

-- Values -----------------------------------------------------------------

data Value a where
  Value :: Show a => !a -> Value a

deriving instance Eq a => Eq (Value a)

deriving instance Ord a => Ord (Value a)

instance Show (Value a) where
  showsPrec :: Int -> Value a -> ShowS
showsPrec Int
p (Value a
a) = Int -> a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
p a
a

unValue :: Value a -> a
unValue :: forall a. Value a -> a
unValue (Value a
v) = a
v

-- Cruft ------------------------------------------------------------------

data Evidence c where
  Evidence :: c => Evidence c

instance Typeable c => Show (Evidence c) where
  show :: Evidence c -> String
show Evidence c
_ = String
"Evidence@(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ forall {k} (t :: k). Typeable t => String
forall (t :: Constraint). Typeable t => String
showType @c String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

unionWithMaybe :: (a -> a -> a) -> Maybe a -> Maybe a -> Maybe a
unionWithMaybe :: forall a. (a -> a -> a) -> Maybe a -> Maybe a -> Maybe a
unionWithMaybe a -> a -> a
f Maybe a
ma Maybe a
ma' = (a -> a -> a
f (a -> a -> a) -> Maybe a -> Maybe (a -> a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe a
ma Maybe (a -> a) -> Maybe a -> Maybe a
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe a
ma') Maybe a -> Maybe a -> Maybe a
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe a
ma Maybe a -> Maybe a -> Maybe a
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe a
ma'

instance Arbitrary a => Arbitrary (NonEmpty a) where
  arbitrary :: Gen (NonEmpty a)
arbitrary = do
    NonEmpty [a]
xs <- Gen (NonEmptyList a)
forall a. Arbitrary a => Gen a
arbitrary
    NonEmpty a -> Gen (NonEmpty a)
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> NonEmpty a
forall a. HasCallStack => [a] -> NonEmpty a
NE.fromList [a]
xs)