Computed.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}

module Data.Monoid.Computed where

import Data.Proxy
import GHC.Generics
import NoThunks.Class
import System.IO.Unsafe (unsafePerformIO)
import Debug.Trace (trace)
import Data.Typeable (Typeable)

-- A newtype wrapper for using with `DeriveVia`.
--
-- A record
-- ```
-- data Rec = Rec { field1 :: a, field2 :: a }
--   deriving Generic
--   deriving Semigroup via (Computed Rec)
-- ```
-- will have an `Evaluated` instance which construct `Rec` such that
-- evaluated fields will be preferred over unevaluated ones (only WHNF is
-- checked).
newtype Computed a = Computed { getComputed :: a }

class Evaluated a where
    -- | Property: if both inputs are equal then the output is equal to them as
    -- well.
    computed :: a -> a -> a

instance (Generic a, GComputed (Rep a)) => Evaluated (Computed a) where
    Computed a `computed` Computed a' = Computed (to (r `gComputed` r'))
      where
        !r  = from a
        !r' = from a'

class GComputed f where
    gComputed :: f a -> f a -> f a

instance GComputed f => GComputed (C1 c f) where
    gComputed (M1 fp) (M1 fp') = M1 (gComputed fp fp')

instance GComputed f => GComputed (D1 c f) where
    gComputed (M1 fp) (M1 fp') = M1 (gComputed fp fp')

instance GComputed f => GComputed (S1 ('MetaSel 'Nothing su ss ds) f) where
    gComputed (M1 fp) (M1 fp') = M1 (gComputed fp fp')

instance GWNoThunks '[] f => GComputed (S1 ('MetaSel ('Just fieldName) su ss ds) f) where
    gComputed a@(M1 fp) a' =
      case unsafePerformIO (gwNoThunks (Proxy @'[]) [] fp) of
        Nothing -> a
        Just _  -> a'

instance (GComputed f, GComputed g) => GComputed (f :*: g) where
  gComputed (a :*: b) (a' :*: b') = a `gComputed` a' :*: b `gComputed` b' 

instance (Typeable c) => GComputed (K1 i c) where
    gComputed a@(K1 c) a' =
      case unsafeNoThunks (OnlyCheckWhnf c) of
        Nothing -> a
        Just _  -> a'

instance GComputed U1 where
    gComputed a _ = a

instance GComputed V1 where
    gComputed a _ = a

data Rec = Rec { field1 :: Int, field2 :: Int }
  deriving (Show, Generic)
  deriving Evaluated via (Computed Rec)

instance NoThunks Rec where

x, y, z, y', z' :: Int
x  = 2
y  = trace "y"  (1 + 1)
y' = trace "y'" (1 + 1)
z  = trace "z"  (1 + 1)
z' = trace "z'" (1 + 1)

rc0, rc1, rc2, rc3, res0, res1 :: Rec
rc0 = x `seq` Rec x y
rc1 = x `seq` Rec z x
rc2 = Rec y' z'
rc3 = Rec z' y'
res0 = rc0 `computed` rc1 -- no trace messages
res1 = rc2 `computed` rc3 -- `y` and `z` are traced

-- note: it's not easy to test this with `NoThunks`, since the constructed
-- record will have thunks, even the constructed fields will be thunks which
-- link to the evaluated code.