{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}

module Constrained.Env where

import Constrained.Core
import Constrained.GenT
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Typeable
import Prettyprinter

-- | Typed environments for mapping `Var a` to `a`
newtype Env = Env {Env -> Map EnvKey EnvValue
unEnv :: Map EnvKey EnvValue}
  deriving newtype (NonEmpty Env -> Env
Env -> Env -> Env
forall b. Integral b => b -> Env -> Env
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: forall b. Integral b => b -> Env -> Env
$cstimes :: forall b. Integral b => b -> Env -> Env
sconcat :: NonEmpty Env -> Env
$csconcat :: NonEmpty Env -> Env
<> :: Env -> Env -> Env
$c<> :: Env -> Env -> Env
Semigroup, Semigroup Env
Env
[Env] -> Env
Env -> Env -> Env
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [Env] -> Env
$cmconcat :: [Env] -> Env
mappend :: Env -> Env -> Env
$cmappend :: Env -> Env -> Env
mempty :: Env
$cmempty :: Env
Monoid)
  deriving stock (Int -> Env -> ShowS
[Env] -> ShowS
Env -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Env] -> ShowS
$cshowList :: [Env] -> ShowS
show :: Env -> String
$cshow :: Env -> String
showsPrec :: Int -> Env -> ShowS
$cshowsPrec :: Int -> Env -> ShowS
Show)

data EnvValue where
  EnvValue :: (Typeable a, Show a) => !a -> EnvValue

deriving instance Show EnvValue

data EnvKey where
  EnvKey :: !(Var a) -> EnvKey

instance Eq EnvKey where
  EnvKey Var a
v == :: EnvKey -> EnvKey -> Bool
== EnvKey Var a
v' = forall a. Var a -> Int
nameOf Var a
v forall a. Eq a => a -> a -> Bool
== forall a. Var a -> Int
nameOf Var a
v'

instance Ord EnvKey where
  compare :: EnvKey -> EnvKey -> Ordering
compare (EnvKey Var a
v) (EnvKey Var a
v') = forall a. Ord a => a -> a -> Ordering
compare (forall a. Var a -> Int
nameOf Var a
v) (forall a. Var a -> Int
nameOf Var a
v')

instance Show EnvKey where
  show :: EnvKey -> String
show (EnvKey Var a
var) = forall a. Show a => a -> String
show Var a
var

extendEnv :: (Typeable a, Show a) => Var a -> a -> Env -> Env
extendEnv :: forall a. (Typeable a, Show a) => Var a -> a -> Env -> Env
extendEnv Var a
v a
a (Env Map EnvKey EnvValue
m) = Map EnvKey EnvValue -> Env
Env forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (forall a. Var a -> EnvKey
EnvKey Var a
v) (forall a. (Typeable a, Show a) => a -> EnvValue
EnvValue a
a) Map EnvKey EnvValue
m

removeVar :: Var a -> Env -> Env
removeVar :: forall a. Var a -> Env -> Env
removeVar Var a
v (Env Map EnvKey EnvValue
m) = Map EnvKey EnvValue -> Env
Env forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (forall a. Var a -> EnvKey
EnvKey Var a
v) Map EnvKey EnvValue
m

singletonEnv :: (Typeable a, Show a) => Var a -> a -> Env
singletonEnv :: forall a. (Typeable a, Show a) => Var a -> a -> Env
singletonEnv Var a
v a
a = Map EnvKey EnvValue -> Env
Env forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
Map.singleton (forall a. Var a -> EnvKey
EnvKey Var a
v) (forall a. (Typeable a, Show a) => a -> EnvValue
EnvValue a
a)

lookupEnv :: Typeable a => Env -> Var a -> Maybe a
lookupEnv :: forall a. Typeable a => Env -> Var a -> Maybe a
lookupEnv (Env Map EnvKey EnvValue
m) Var a
v = do
  EnvValue a
val <- forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (forall a. Var a -> EnvKey
EnvKey Var a
v) Map EnvKey EnvValue
m
  forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast a
val

findEnv :: (Typeable a, MonadGenError m) => Env -> Var a -> m a
findEnv :: forall a (m :: * -> *).
(Typeable a, MonadGenError m) =>
Env -> Var a -> m a
findEnv Env
env Var a
var = do
  case forall a. Typeable a => Env -> Var a -> Maybe a
lookupEnv Env
env Var a
var of
    Just a
a -> forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
    Maybe a
Nothing -> forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
genError (forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
"Couldn't find " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Var a
var forall a. [a] -> [a] -> [a]
++ String
" in " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Env
env))

instance Pretty EnvValue where
  pretty :: forall ann. EnvValue -> Doc ann
pretty (EnvValue a
x) = forall a ann. Pretty a => a -> Doc ann
pretty forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take Int
80 (forall a. Show a => a -> String
show a
x)

instance Pretty EnvKey where
  pretty :: forall ann. EnvKey -> Doc ann
pretty (EnvKey Var a
x) = forall a ann. Show a => a -> Doc ann
viaShow Var a
x

instance Pretty Env where
  pretty :: forall ann. Env -> Doc ann
pretty (Env Map EnvKey EnvValue
m) = forall ann. [Doc ann] -> Doc ann
vsep (Doc ann
"Env" forall a. a -> [a] -> [a]
: (forall a b. (a -> b) -> [a] -> [b]
map forall {a} {a} {ann}. (Pretty a, Pretty a) => (a, a) -> Doc ann
f (forall k a. Map k a -> [(k, a)]
Map.toList Map EnvKey EnvValue
m)))
    where
      f :: (a, a) -> Doc ann
f (a
k, a
v) = forall ann. [Doc ann] -> Doc ann
hsep [forall a ann. Pretty a => a -> Doc ann
pretty a
k, Doc ann
"->", forall a ann. Pretty a => a -> Doc ann
pretty a
v]