From 9d981ee013479197e5dc57b92bae94d915dd402b Mon Sep 17 00:00:00 2001
From: David Feuer <David.Feuer@gmail.com>
Date: Wed, 21 Dec 2022 03:27:23 -0500
Subject: [PATCH] Add strict traversal operations

* Add `strictly` to turn a lazy (standard) traversal into a strict one
  that forces targets before installing them.
* Add `over'`, `iover'`, `modifying'`, `imodifying'`, and corresponding
  operators `%!~`, `%!@~`, `%!=`, and `%!@=`.
* Adjust documentation.

Closes #1016
---
 lens.cabal                          |   2 +
 src/Control/Lens/Combinators.hs     |   4 +
 src/Control/Lens/Internal/BoxT.hs   |  37 ++++++
 src/Control/Lens/Internal/Setter.hs |   9 ++
 src/Control/Lens/Setter.hs          |  14 ++
 src/Control/Lens/Traversal.hs       | 197 +++++++++++++++++++++++++++-
 6 files changed, 262 insertions(+), 1 deletion(-)
 create mode 100644 src/Control/Lens/Internal/BoxT.hs

diff --git a/lens.cabal b/lens.cabal
index 6acf6115f..519ea4109 100644
--- a/lens.cabal
+++ b/lens.cabal
@@ -190,6 +190,7 @@ library
     indexed-traversable-instances >= 0.1      && < 0.2,
     kan-extensions                >= 5        && < 6,
     mtl                           >= 2.2.1    && < 2.4,
+    OneTuple                      >= 0.3      && < 0.4,
     parallel                      >= 3.2.1.0  && < 3.3,
     profunctors                   >= 5.5.2    && < 6,
     reflection                    >= 2.1      && < 3,
@@ -223,6 +224,7 @@ library
     Control.Lens.Indexed
     Control.Lens.Internal
     Control.Lens.Internal.Bazaar
+    Control.Lens.Internal.BoxT
     Control.Lens.Internal.ByteString
     Control.Lens.Internal.Context
     Control.Lens.Internal.CTypes
diff --git a/src/Control/Lens/Combinators.hs b/src/Control/Lens/Combinators.hs
index 4174ee661..502680e1c 100644
--- a/src/Control/Lens/Combinators.hs
+++ b/src/Control/Lens/Combinators.hs
@@ -107,6 +107,7 @@ import Control.Lens hiding
   , (...)
   , (#)
   , (%~)
+  , (%!~)
   , (.~)
   , (?~)
   , (<.~)
@@ -124,6 +125,7 @@ import Control.Lens hiding
   , (&&~)
   , (.=)
   , (%=)
+  , (%!=)
   , (?=)
   , (+=)
   , (-=)
@@ -140,7 +142,9 @@ import Control.Lens hiding
   , (<>~)
   , (<>=)
   , (%@~)
+  , (%!@~)
   , (%@=)
+  , (%!@=)
   , (:>)
   , (:<)
   )
diff --git a/src/Control/Lens/Internal/BoxT.hs b/src/Control/Lens/Internal/BoxT.hs
new file mode 100644
index 000000000..d82701699
--- /dev/null
+++ b/src/Control/Lens/Internal/BoxT.hs
@@ -0,0 +1,37 @@
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DeriveTraversable #-}
+module Control.Lens.Internal.BoxT where
+import Control.Applicative
+import Data.Functor.Apply (Apply (..))
+import Data.Functor.Contravariant (Contravariant (..))
+import Data.Tuple.Solo (Solo (..))
+
+-- | A very simple applicative transformer that gives us more control over when
+-- things get forced. Note: this type /should not/ be made an instance of
+-- @Settable@, because then users could accidentally use
+-- 'Control.Lens.Traversable.strictly' with a 'Control.Lens.Setter.Setter',
+-- which will not work at all. There is no way to strictify a @Setter@.
+newtype BoxT f a = BoxT
+  { runBoxT :: f (Solo a) }
+  deriving (Functor, Foldable, Traversable)
+
+-- The Contravariant instance allows `strictly` to be used on a getter or fold.
+-- It's not at all obvious that this is *useful* (since `strictly` doesn't
+-- change these at all), but it's also not obviously *harmful*.
+instance Contravariant f => Contravariant (BoxT f) where
+  contramap f (BoxT m) = BoxT $ contramap (fmap f) m
+  {-# INLINE contramap #-}
+instance Apply f => Apply (BoxT f) where
+  liftF2 f (BoxT m) (BoxT n) = BoxT (liftF2 (liftA2 f) m n)
+  {-# INLINE liftF2 #-}
+instance Applicative f => Applicative (BoxT f) where
+  pure = BoxT . pure . Solo
+  {-# INLINE pure #-}
+  BoxT m <*> BoxT n = BoxT (liftA2 (<*>) m n)
+  {-# INLINE (<*>) #-}
+#if MIN_VERSION_base(4,10,0)
+  liftA2 f (BoxT m) (BoxT n) = BoxT (liftA2 (liftA2 f) m n)
+  {-# INLINE liftA2 #-}
+#endif
+  -- Caution: We *can't* implement *> or <* in terms of the underlying *> and
+  -- <*. We need to force the Solos, not discard them.
diff --git a/src/Control/Lens/Internal/Setter.hs b/src/Control/Lens/Internal/Setter.hs
index 48569ff58..cb20929ab 100644
--- a/src/Control/Lens/Internal/Setter.hs
+++ b/src/Control/Lens/Internal/Setter.hs
@@ -49,6 +49,15 @@ instance Settable Identity where
   taintedDot = (Identity #.)
   {-# INLINE taintedDot #-}
 
+-- CAUTION: While Data.Tuple.Solo may *look* a lot like Identity, and while we
+-- *could* give it a Settable instance, we probably do not want to do so. In
+-- particular, if we did, then Control.Lens.Traversal.over' would "work" with
+-- Setters. But ... it wouldn't *actually* work; the mapping would end up being
+-- lazy when it's supposed to be strict. Similarly, the BoxT applicative
+-- transformer must not be made Settable, because that would cause a similarly
+-- confusing problem with Control.Lens.Traversal.strictly. There is not, as
+-- yet, any compelling reason to write such an instance, so let's not.
+
 -- | 'Control.Lens.Fold.backwards'
 instance Settable f => Settable (Backwards f) where
   untainted = untaintedDot forwards
diff --git a/src/Control/Lens/Setter.hs b/src/Control/Lens/Setter.hs
index b2168cdf4..231a56090 100644
--- a/src/Control/Lens/Setter.hs
+++ b/src/Control/Lens/Setter.hs
@@ -342,6 +342,13 @@ cloneIndexedSetter l pafb = taintedDot (runIdentity #. l (Indexed $ \i -> Identi
 -- >>> over _1 show (10,20)
 -- ("10",20)
 --
+--
+-- Like 'fmap', @over@ is normally lazy in the result(s) of calling the
+-- function, which can cause space leaks in lazy fields, or when using
+-- 'Control.Lens.At.ix' for value-lazy structures like 'Data.Sequence.Seq',
+-- 'Data.Map.Map', 'Data.IntMap.IntMap', or 'Data.Array.Array'.  For a strict
+-- version, see `Control.Lens.Traversal.iover'`.
+--
 -- @
 -- 'over' :: 'Setter' s t a b -> (a -> b) -> s -> t
 -- 'over' :: 'ASetter' s t a b -> (a -> b) -> s -> t
@@ -1169,6 +1176,13 @@ ilocally l f = Reader.local (iover l f)
 -- 'iover' l ≡ 'over' l '.' 'Indexed'
 -- @
 --
+-- Like 'Data.Functor.WithIndex.imap', @iover@ is normally lazy in the
+-- result(s) of calling the function, which can cause space leaks in lazy
+-- fields, or when using 'Control.Lens.At.ix' for value-lazy structures like
+-- 'Data.Sequence.Seq', 'Data.Map.Map', 'Data.IntMap.IntMap', or
+-- 'Data.Array.Array'.  For a strict version, see
+-- `Control.Lens.Traversal.iover'`.
+--
 -- @
 -- 'iover' :: 'IndexedSetter' i s t a b    -> (i -> a -> b) -> s -> t
 -- 'iover' :: 'IndexedLens' i s t a b      -> (i -> a -> b) -> s -> t
diff --git a/src/Control/Lens/Traversal.hs b/src/Control/Lens/Traversal.hs
index 1b1d57e6d..7bed4e32d 100644
--- a/src/Control/Lens/Traversal.hs
+++ b/src/Control/Lens/Traversal.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE Rank2Types #-}
 {-# LANGUAGE FlexibleContexts #-}
@@ -6,6 +7,7 @@
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE Trustworthy #-}
+{-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE ConstraintKinds #-}
 
 #include "lens-common.h"
@@ -113,6 +115,17 @@ module Control.Lens.Traversal
   , imapAccumROf
   , imapAccumLOf
 
+  -- ** Strict traversals
+  , over'
+  , (%!~)
+  , iover'
+  , (%!@~)
+  , modifying'
+  , (%!=)
+  , imodifying'
+  , (%!@=)
+  , strictly
+
   -- * Reflection
   , traverseBy
   , traverseByOf
@@ -137,6 +150,7 @@ import Control.Comonad
 import Control.Lens.Fold
 import Control.Lens.Getter (Getting, IndexedGetting, getting)
 import Control.Lens.Internal.Bazaar
+import Control.Lens.Internal.BoxT
 import Control.Lens.Internal.Context
 import Control.Lens.Internal.Fold
 import Control.Lens.Internal.Indexed
@@ -145,6 +159,8 @@ import Control.Lens.Lens
 import Control.Lens.Setter (ASetter, AnIndexedSetter, isets, sets)
 import Control.Lens.Type
 import Control.Monad.Trans.State.Lazy
+import Control.Monad.State.Class (MonadState)
+import qualified Control.Monad.State.Class as MonadState
 import Data.Bitraversable
 import Data.CallStack
 import Data.Functor.Apply
@@ -162,10 +178,11 @@ import Data.Reflection
 import Data.Semigroup.Traversable
 import Data.Semigroup.Bitraversable
 import Data.Tuple (swap)
+import Data.Tuple.Solo (Solo (..), getSolo)
 import GHC.Magic (inline)
 
 -- $setup
--- >>> :set -XNoOverloadedStrings -XFlexibleContexts
+-- >>> :set -XNoOverloadedStrings -XFlexibleContexts -XRankNTypes
 -- >>> import Data.Char (toUpper)
 -- >>> import Control.Applicative
 -- >>> import Control.Lens
@@ -183,6 +200,9 @@ import GHC.Magic (inline)
 -- >>> let firstAndThird :: Traversal (a, x, a) (b, x, b) a b; firstAndThird = traversal go where { go :: Applicative f => (a -> f b) -> (a, x, a) -> f (b, x, b); go focus (a, x, a') = liftA3 (,,) (focus a) (pure x) (focus a') }
 -- >>> let selectNested :: Traversal (x, [a]) (x, [b]) a b; selectNested = traversal go where { go :: Applicative f => (a -> f b) -> (x, [a]) -> f (x, [b]); go focus (x, as) = liftA2 (,) (pure x) (traverse focus as) }
 
+infixr 4 %!~, %!@~
+infix  4 %!=, %!@=
+
 ------------------------------------------------------------------------------
 -- Traversals
 ------------------------------------------------------------------------------
@@ -1466,3 +1486,178 @@ traverseByOf l pur app f = reifyApplicative pur app (l (ReflectedApplicative #.
 -- @
 sequenceByOf :: Traversal s t (f b) b -> (forall x. x -> f x) -> (forall x y. f (x -> y) -> f x -> f y) -> s -> f t
 sequenceByOf l pur app = reifyApplicative pur app (l ReflectedApplicative)
+
+-- Note: Solo wrapping
+--
+-- We use Solo for strict application of (indexed) setters.
+--
+-- Credit for this idea goes to Eric Mertens; see
+-- <https://github.com/glguy/irc-core/commit/2d5fc45b05f1>. It was reinvented
+-- independently by David Feuer, who realized that an applicative transformer
+-- version could be used to implement `strictly`.
+--
+-- Using Solo rather than Identity allows us, when applying a traversal to a
+-- structure, to evaluate only the parts that we modify. If an optic focuses on
+-- multiple targets, the Applicative instance of Solo (combined with applying
+-- the Solo data constructor strictly) makes sure that we force evaluation of
+-- all of them, but we leave anything else alone.
+
+-- | A version of 'Control.Lens.Setter.over' that forces the result(s) of
+-- applying the function. This can prevent space leaks when modifying lazy
+-- fields. See also 'strictly'.
+--
+-- @
+-- over' :: 'Lens' s t a b -> (a -> b) -> s -> t
+-- over' :: 'Traversal' s t a b -> (a -> b) -> s -> t
+-- @
+--
+-- >>> length $ over traverse id [undefined, undefined]
+-- 2
+--
+-- >>> over' traverse id [1, undefined :: Int]
+-- *** Exception: Prelude.undefined
+-- ...
+over' :: LensLike Solo s t a b -> (a -> b) -> s -> t
+-- See [Note: Solo wrapping]
+over' l f = getSolo . l (\old -> Solo $! f old)
+{-# INLINE over' #-}
+
+-- | Traverse targets strictly. This is the operator version of 'over''.
+(%!~) :: LensLike Solo s t a b -> (a -> b) -> s -> t
+(%!~) = over'
+{-# INLINE (%!~) #-}
+
+-- $
+-- >>> :{
+-- let lover' :: Lens s t a b -> (a -> b) -> s -> t
+--     lover' l = over' l
+--     tover' :: Traversal s t a b -> (a -> b) -> s -> t
+--     tover' l = over' l
+-- :}
+--
+-- >>> :{
+-- let sover' :: Setter s t a b -> (a -> b) -> s -> t
+--     sover' l = over' l
+-- :}
+-- ...
+-- ...error...
+-- ...
+
+-- | A version of 'Control.Lens.Setter.iover' that forces the result(s) of
+-- applying the function. Alternatively, an indexed version of `over'`.
+-- See also 'strictly'.
+--
+-- @
+-- iover' :: IndexedLens i s t a b -> (i -> a -> b) -> s -> t
+-- iover' :: IndexedTraversal i s t a b -> (i -> a -> b) -> s -> t
+-- @
+iover' :: Over (Indexed i) Solo s t a b -> (i -> a -> b) -> s -> t
+-- See [Note: Solo wrapping]
+iover' l f = getSolo . l (Indexed $ \i a -> Solo $! f i a)
+{-# INLINE iover' #-}
+
+-- | Traverse targets strictly with indices. This is the operator version of
+-- 'iover''.
+(%!@~) :: Over (Indexed i) Solo s t a b -> (i -> a -> b) -> s -> t
+(%!@~) = iover'
+{-# INLINE (%!@~) #-}
+
+-- | Modify the state strictly. This is stricter than
+-- @Control.Lens.Setter.modifying@ in two ways: it forces the new value of the
+-- state, and it forces the new value of the target within the state.
+modifying' :: MonadState s m => LensLike Solo s s a b -> (a -> b) -> m ()
+-- See [Note: Solo wrapping]
+modifying' l f = do
+  s <- MonadState.get
+  let !(Solo !t) = l (\old -> Solo $! f old) s
+  MonadState.put t
+{-# INLINE modifying' #-}
+
+-- | Modify the state strictly. This is an operator version of
+-- 'modifying''.
+(%!=) :: MonadState s m => LensLike Solo s s a b -> (a -> b) -> m ()
+(%!=) = modifying'
+{-# INLINE (%!=) #-}
+
+-- | Modify the state strictly with an index. This is stricter than
+-- @Control.Lens.Setter.imodifying@ in two ways: it forces the new value of the
+-- state, and it forces the new value of the target within the state.
+imodifying' :: MonadState s m => Over (Indexed i) Solo s s a b -> (i -> a -> b) -> m ()
+-- See [Note: Solo wrapping]
+imodifying' l f = do
+  s <- MonadState.get
+  let !(Solo !t) = l (Indexed $ \i old -> Solo $! f i old) s
+  MonadState.put t
+{-# INLINE imodifying' #-}
+
+-- | Modify the state strictly. This is an operator version of
+-- 'imodifying''.
+(%!@=) :: MonadState s m => Over (Indexed i) Solo s s a b -> (i -> a -> b) -> m ()
+(%!@=) = imodifying'
+{-# INLINE (%!@=) #-}
+
+-- $
+-- >>> :{
+-- let liover' :: IndexedLens i s t a b -> (i -> a -> b) -> s -> t
+--     liover' l = iover' l
+--     tiover' :: IndexedTraversal i s t a b -> (i -> a -> b) -> s -> t
+--     tiover' l = iover' l
+-- :}
+
+-- | Use an optic /strictly/. @strictly l f s@ will force the results of /all/
+-- the targets of @l@ when a new outer value is forced.
+--
+-- @strictly@ does not affect folds or getters in any way, as they don't produce
+-- new outer values.
+--
+-- Note that producing an optic using 'strictly' will not necessarily produce
+-- one as efficient as what could be written by hand, although it will do so in
+-- simple enough situations. Efficiency issues are most likely when working
+-- over a large structure in a functor other than the usual 'Identity'.
+--
+-- @
+-- 'over'' l = 'Control.Lens.Setter.over' (strictly l)
+-- 'iover'' l = 'Control.Lens.Setter.iover' (strictly l)
+-- @
+--
+-- @
+-- strictly :: 'Traversal' s t a b -> 'Traversal' s t a b
+-- strictly :: 'IndexedTraversal' i s t a b -> 'IndexedTraversal' i s t a b
+-- @
+strictly :: (Functor f, Profunctor p, Profunctor q) => Optical p q (BoxT f) s t a b -> Optical p q f s t a b
+-- See [Note: Solo wrapping]
+strictly l f = rmap (fmap getSolo .# runBoxT) $ l (rmap (BoxT #. fmap (Solo $!)) f)
+{-# INLINE strictly #-}
+
+{-
+-- If the ambient functor is either a Traversable or a Monad, then we can get
+-- rid of the Solo boxes ourselves:
+
+strictlyT :: (Traversable f, Profunctor p, Profunctor q) => Optical p q (BoxT f) s t a b -> Optical p q f s t a b
+strictlyT l f = rmap (getSolo . sequenceA .# runBoxT) $ l (rmap (BoxT #. fmap (Solo $!)) f)
+
+strictlyM :: (Monad f, Profunctor p, Profunctor q) => Optical p q (BoxT f) s t a b -> Optical p q f s t a b
+strictlyM l f = rmap ((>>= \(Solo r) -> pure r) .# runBoxT) $ l (rmap (BoxT #. fmap (Solo $!)) f)
+-}
+
+-- $
+-- >>> :{
+-- let tstrictly :: Traversal s t a b -> Traversal s t a b
+--     tstrictly l = strictly l
+--     itstrictly :: AnIndexedTraversal i s t a b -> IndexedTraversal i s t a b
+--     itstrictly l = strictly (cloneIndexedTraversal l)
+--     lstrictly :: Lens s t a b -> Lens s t a b
+--     lstrictly l = strictly l
+--     ilstrictly :: AnIndexedLens i s t a b -> IndexedLens i s t a b
+--     ilstrictly l = strictly (cloneIndexedLens l)
+--     fstrictly :: Fold s a -> Fold s a
+--     fstrictly l = strictly l
+-- :}
+--
+-- >>> :{
+-- let sstrictly :: Setter s t a b -> Setter s t a b
+--     sstrictly l = strictly l
+-- :}
+-- ...
+-- ...Settable ...BoxT...
+-- ...