Refine.hs

{-# LANGUAGE AllowAmbiguousTypes #-} -- for predicate reification
{-# LANGUAGE UndecidableInstances #-} -- for KnownPredicateName in Arbitrary

module Rerefined.Refine
  (
  -- * @Refined@
    type Refined
  , refine
  , unrefine
  , reifyPredicate
  , unsafeRefine
  , unsafeRerefine

  -- * @Refined1@
  , type Refined1
  , refine1
  , unrefine1
  , reifyPredicate1
  , squashRefined1
  , unsafeRefine1
  , unsafeRerefine1

  -- * Errors
  , type RefineFailure
  , prettyRefineFailure
  , prettyRefineFailure'
  ) where

import Rerefined.Predicate
import Language.Haskell.TH.Syntax ( Lift )
import GHC.Exts ( proxy# )
import Data.Text ( Text )
import Data.Text.Builder.Linear qualified as TBL

import Test.QuickCheck ( Arbitrary(arbitrary) )
import Test.QuickCheck.Gen qualified as QC

-- | @a@ refined with predicate @p@.
newtype Refined p a = Refined a
    deriving stock (Lift, Eq, Ord, Show)

-- | Strip the refinement from a 'Refined'.
--
-- This is kept as a separate function for prettier @'Show' 'Refined'@ output.
unrefine :: Refined p a -> a
unrefine (Refined a) = a

-- | @f a@ refined with predicate @p@.
--
-- We may derive legal 'Functor', 'Traversable' instances for this as
-- 'Rerefined.Predicate.Refine1' guarantees that the predicate only applies to
-- the functor structure. That is, you _may_ alter a 'Refined1' without
-- re-asserting its predicate, provided your changes are made without altering
-- the structure/shape of @f@ (e.g. 'fmap', 'traverse').
newtype Refined1 p f a = Refined1 (f a)
    deriving stock (Functor, Foldable, Traversable, Lift, Eq, Ord, Show)

-- | Strip the refinement from a 'Refined1'.
--
-- This is kept as a separate function for prettier @'Show' 'Refined1'@ output.
unrefine1 :: Refined1 p f a -> f a
unrefine1 (Refined1 fa) = fa

-- | Squash a 'Refined1' into a 'Refined'. Essentially forget the @f@.
squashRefined1 :: Refined1 p f a -> Refined p (f a)
squashRefined1 = Refined . unrefine1

-- | Refine @a@ with predicate @p@.
refine
    :: forall p a. Refine p a
    => a -> Either RefineFailure (Refined p a)
refine a =
    case validate (proxy# @p) a of
      Nothing -> Right (Refined a)
      Just e  -> Left e

-- | Construct a 'Refined' without validating the predicate @p@.
--
-- Unsafe. Use only when you can manually prove that the predicate holds.
unsafeRefine :: a -> Refined p a
unsafeRefine = Refined

-- | Replace a 'Refined''s predicate without validating the new prdicate @pNew@.
--
-- Unsafe. Use only when you can manually prove that the new predicate holds.
unsafeRerefine :: forall pNew pOld a. Refined pOld a -> Refined pNew a
unsafeRerefine = Refined . unrefine

-- | Reify a predicate.
reifyPredicate :: forall p a. Refine p a => a -> Bool
reifyPredicate a = case refine @p a of Right{} -> True; Left{} -> False

-- | Refine @f a@ with functor predicate @p@.
refine1
    :: forall p f a. Refine1 p f
    => f a -> Either RefineFailure (Refined1 p f a)
refine1 fa =
    case validate1 (proxy# @p) fa of
      Nothing -> Right (Refined1 fa)
      Just e  -> Left e

-- | Construct a 'Refined1' without validating the predicate @p@.
--
-- Unsafe. Use only when you can manually prove that the predicate holds.
unsafeRefine1 :: f a -> Refined1 p f a
unsafeRefine1 = Refined1

-- | Replace a 'Refined1''s predicate without validating the new prdicate
--   @pNew@.
--
-- Unsafe. Use only when you can manually prove that the new predicate holds.
unsafeRerefine1 :: forall pNew pOld f a. Refined1 pOld f a -> Refined1 pNew f a
unsafeRerefine1 = Refined1 . unrefine1

-- | Reify a functor predicate.
reifyPredicate1 :: forall p f a. Refine1 p f => f a -> Bool
reifyPredicate1 fa = case refine1 @p fa of Right{} -> True; Left{} -> False

{- TODO
* got an extra \n at start oops
* make it look better lol
* make tail-call recursive? need to ferry more indents around though
-}
prettyRefineFailure :: RefineFailure -> Text
prettyRefineFailure = TBL.runBuilder . go (0 :: Int) . (\e -> [e])
  where
    go n = \case
      []     -> mempty
      (e:es) ->
        let bPred   = TBL.fromChar '\n' <> indent n <> refineFailurePredicate e
            bDetail = TBL.fromChar '\n' <> indent (n+2) <> refineFailureDetail e
         in bPred <> bDetail <> go (n+2) (refineFailureInner e) <> go n es
    indent = \case
      0 -> mempty
      n -> TBL.fromChar ' ' <> indent (n-1)

-- TODO this requires switching inner errors so that last is first lol x)
-- also we only <> right. maybe that's useful for perf
-- to remove newline at start we need to start in a special mode so we know not
-- to add the first newline at bPred. but I cba it will be messier and I want to
-- replace this ASAP when someone comes up with a good pretty error format
prettyRefineFailure' :: RefineFailure -> Text
prettyRefineFailure' = \e -> TBL.runBuilder $ go mempty [(0 :: Int, e)]
  where
    go b = \case
      []            -> b
      ((n, e) : es) ->
        let bPred   = TBL.fromChar '\n' <> indent n <> refineFailurePredicate e
            bDetail = TBL.fromChar '\n' <> indent (n+2) <> refineFailureDetail e
            b'      = b <> bPred <> bDetail
         in go b' (idk (n+2) es (refineFailureInner e))
    indent = \case
      0 -> mempty
      n -> TBL.fromChar ' ' <> indent (n-1)
    idk n rs = \case
      []   -> rs
      l:ls -> idk n ((n, l):rs) ls

-- | Generate a refined term by generating an unrefined term and asserting the
--   predicate.
--
-- Will runtime error if it fails to find a satisfying term (based on size).
instance (Arbitrary a, Refine p a, KnownPredicateName p)
  => Arbitrary (Refined p a) where
    arbitrary = QC.suchThatMaybe arbitrary (reifyPredicate @p) >>= \case
      Just a  -> pure $ Refined a
      Nothing -> error $
           "rerefined: couldn't generate a value satisfying predicate: "
        <> predicateName @p

-- | Generate a refined term by generating an unrefined term and asserting the
--   functor predicate.
--
-- Will runtime error if it fails to find a satisfying term (based on size).
instance (Arbitrary (f a), Refine1 p f, KnownPredicateName p)
  => Arbitrary (Refined1 p f a) where
    arbitrary = QC.suchThatMaybe arbitrary (reifyPredicate1 @p) >>= \case
      Just fa -> pure $ Refined1 fa
      Nothing -> error $
           "rerefined: couldn't generate a value satisfying predicate: "
        <> predicateName @p