From 9fcd3c882927a421359039dea482c42059b5565a Mon Sep 17 00:00:00 2001
From: Kiana Sheibani <kiana.a.sheibani@gmail.com>
Date: Fri, 14 Apr 2023 13:16:08 -0400
Subject: [PATCH] Add more utility functions

---
 src/Control/Applicative/Indexing.idr | 33 +++++++++++++
 src/Control/Lens/Equality.idr        |  2 +-
 src/Control/Lens/Fold.idr            | 63 +++++++++++++++++++++--
 src/Control/Lens/Getter.idr          |  5 ++
 src/Control/Lens/Lens.idr            |  8 +++
 src/Control/Lens/Optional.idr        | 15 +++---
 src/Control/Lens/Prism.idr           | 25 ++++++++++
 src/Control/Lens/Traversal.idr       | 74 ++++++++++++++++++++++++++--
 8 files changed, 208 insertions(+), 17 deletions(-)
 create mode 100644 src/Control/Applicative/Indexing.idr

diff --git a/src/Control/Applicative/Indexing.idr b/src/Control/Applicative/Indexing.idr
new file mode 100644
index 0000000..ba51795
--- /dev/null
+++ b/src/Control/Applicative/Indexing.idr
@@ -0,0 +1,33 @@
+module Control.Applicative.Indexing
+
+import Data.Contravariant
+
+%default total
+
+
+public export
+record Indexing {0 k : Type} f (a : k) where
+  constructor MkIndexing
+  runIndexing : Nat -> (Nat, f a)
+
+
+public export
+Functor f => Functor (Indexing f) where
+  map f (MkIndexing g) = MkIndexing (mapSnd (map f) . g)
+
+public export
+Applicative f => Applicative (Indexing f) where
+  pure x = MkIndexing $ \i => (i, pure x)
+  MkIndexing mf <*> MkIndexing ma = MkIndexing $ \i =>
+    let (j, ff) = mf i
+        (k, fa) = ma j
+    in (k, ff <*> fa)
+
+public export
+Contravariant f => Contravariant (Indexing f) where
+  contramap f (MkIndexing g) = MkIndexing (mapSnd (contramap f) . g)
+
+
+public export
+indexing : ((a -> Indexing f b) -> s -> Indexing f t) -> (Nat -> a -> f b) -> s -> f t
+indexing l fn s = snd $ runIndexing {f} (l (\x => MkIndexing $ \i => (S i, fn i x)) s) 0
diff --git a/src/Control/Lens/Equality.idr b/src/Control/Lens/Equality.idr
index 6bd66c8..683e0a1 100644
--- a/src/Control/Lens/Equality.idr
+++ b/src/Control/Lens/Equality.idr
@@ -66,7 +66,7 @@ substEq : forall p. Equality s t a b -> p a b a b -> p s t a b
 substEq {p} l = l {p = \x,y => p x y a b}
 
 
-||| A trivial `Simple Equality`.
+||| The trivial `Simple Equality`.
 ||| Composing this optic with any other can force it to be a `Simple` optic.
 public export
 simple : Equality' a a
diff --git a/src/Control/Lens/Fold.idr b/src/Control/Lens/Fold.idr
index e4d5261..a23a31b 100644
--- a/src/Control/Lens/Fold.idr
+++ b/src/Control/Lens/Fold.idr
@@ -70,8 +70,8 @@ folding @{_} f @{MkIsFold _} = rphantom . contramapFst f . wander traverse_
 
 ||| Reverse the order of a fold's focuses.
 public export
-backwardsFold : Fold s a -> Fold s a
-backwardsFold l @{MkIsFold _} = rphantom . wander func
+backwards : Fold s a -> Fold s a
+backwards l @{MkIsFold _} = rphantom . wander func
   where
     traversing : Applicative f => Traversing (Forget (f ()))
     traversing =
@@ -82,6 +82,16 @@ backwardsFold l @{MkIsFold _} = rphantom . wander func
     func fn = let _ = traversing in
       forwards . runForget (l $ MkForget (MkBackwards {f} . ignore . fn))
 
+||| Construct a fold that replicates the input n times.
+public export
+replicated : Nat -> Fold a a
+replicated n @{MkIsFold _} = rphantom . wander (\f,x => rep n (f x))
+  where
+    rep : Applicative f => Nat -> f a -> f ()
+    rep Z _ = pure ()
+    rep (S Z) x = ignore x
+    rep (S n@(S _)) x = x *> rep n x
+
 
 ||| Map each focus of an optic to a monoid value and combine them.
 public export
@@ -105,6 +115,11 @@ public export
 concatOf : Monoid m => Fold s m -> s -> m
 concatOf l = foldMapOf l id
 
+||| Fold over the focuses of an optic using Alternative.
+public export
+choiceOf : Alternative f => Fold s (Lazy (f a)) -> s -> f a
+choiceOf = force .: concatOf @{MonoidAlternative}
+
 ||| Evaluate each computation of an optic and discard the results.
 public export
 sequenceOf_ : Applicative f => Fold s (f a) -> s -> f ()
@@ -146,6 +161,32 @@ anyOf : Fold s a -> (a -> Bool) -> s -> Bool
 anyOf = foldMapOf @{Any}
 
 
+||| Return `True` if the element occurs in the focuses of the optic.
+public export
+elemOf : Eq a => Fold s a -> a -> s -> Bool
+elemOf l = allOf l . (==)
+
+||| Calculate the number of focuses of the optic.
+public export
+lengthOf : Fold s a -> s -> Nat
+lengthOf l = foldMapOf @{Additive} l (const 1)
+
+||| Access the first focus value of an optic, returning `Nothing` if there are
+||| no focuses.
+|||
+||| This is identical to `preview`.
+public export
+firstOf : Fold s a -> s -> Maybe a
+firstOf l = foldMapOf l Just
+
+||| Access the last focus value of an optic, returning `Nothing` if there are
+||| no focuses.
+public export
+lastOf : Fold s a -> s -> Maybe a
+lastOf l = foldMapOf @{MkMonoid @{MkSemigroup $ flip (<+>)} neutral} l Just
+
+
+
 ------------------------------------------------------------------------------
 -- Accessing folds
 ------------------------------------------------------------------------------
@@ -166,13 +207,15 @@ hasn't l = allOf l (const False)
 ||| returning `Nothing` if there are no focuses.
 public export
 previews : Fold s a -> (a -> r) -> s -> Maybe r
-previews l f = foldMapOf @{MonoidAlternative} l (Just . f)
+previews l f = foldMapOf l (Just . f)
 
 ||| Access the first focus value of an optic, returning `Nothing` if there are
 ||| no focuses.
+|||
+||| This is an alias for `firstOf`.
 public export
 preview : Fold s a -> s -> Maybe a
-preview l = foldMapOf @{MonoidAlternative} l Just
+preview = firstOf
 
 infixl 8 ^?
 
@@ -185,12 +228,24 @@ public export
 (^?) s l = preview l s
 
 
+||| Convert a `Fold` into an `OptionalFold` that accesses the first focus element.
+|||
+||| For the traversal version of this, see `singular`.
+public export
+pre : Fold s a -> OptionalFold s a
+pre = folding . preview
+
+
+||| Return a list of all focuses of a fold.
 public export
 toListOf : Fold s a -> s -> List a
 toListOf l = foldrOf l (::) []
 
 infixl 8 ^..
 
+||| Return a list of all focuses of a fold.
+|||
+||| This is the operator form of `toListOf`.
 public export
 (^..) : s -> Fold s a -> List a
 (^..) s l = toListOf l s
diff --git a/src/Control/Lens/Getter.idr b/src/Control/Lens/Getter.idr
index 3cc8f2f..e498498 100644
--- a/src/Control/Lens/Getter.idr
+++ b/src/Control/Lens/Getter.idr
@@ -42,6 +42,11 @@ public export
 to : (s -> a) -> Getter s a
 to f @{MkIsGetter _} = lmap f . rphantom
 
+||| Construct a getter that always returns a constant value.
+public export
+like : a -> Getter s a
+like = to . const
+
 
 ||| Access the value of an optic and apply a function to it, returning the result.
 public export
diff --git a/src/Control/Lens/Lens.idr b/src/Control/Lens/Lens.idr
index 43ad30b..76de836 100644
--- a/src/Control/Lens/Lens.idr
+++ b/src/Control/Lens/Lens.idr
@@ -90,6 +90,14 @@ public export
 united : Lens' a ()
 united @{MkIsLens _} = dimap ((),) snd . first
 
+||| Create a lens that operates on pairs from two other lenses.
+public export
+alongside : Lens s t a b -> Lens s' t' a' b' -> Lens (s, s') (t, t') (a, a') (b, b')
+alongside l l' =
+  let (get1, set1) = getLens l
+      (get2, set2) = getLens l'
+  in lens (bimap get1 get2) (uncurry bimap . bimap set1 set2)
+
 
 
 ||| Optimize a composition of optics by fusing their uses of `dimap` into one.
diff --git a/src/Control/Lens/Optional.idr b/src/Control/Lens/Optional.idr
index f9d7265..c9ddb69 100644
--- a/src/Control/Lens/Optional.idr
+++ b/src/Control/Lens/Optional.idr
@@ -64,14 +64,7 @@ optional' prj = optional (\x => maybe (Left x) Right (prj x))
 ||| The trivial optic that has no focuses.
 public export
 ignored : Optional s s a b
-ignored @{MkIsOptional _} = dimap @{fromStrong}
-  (\x => (Left x, const x))
-  (\(e, f) => either id (the (b -> s) f) e)
-  . first . right
-  where
-    -- arbitrary choice of where to pull profunctor instance from
-    fromStrong : Strong p => Profunctor p
-    fromStrong = %search
+ignored = optional' (const Nothing) const
 
 
 ||| Extract projection and setter functions from an optional.
@@ -96,3 +89,9 @@ getOptional l = l @{MkIsOptional (strong,choice)} (Right, const id)
 public export
 withOptional : Optional s t a b -> ((s -> Either t a) -> (s -> b -> t) -> r) -> r
 withOptional l f = uncurry f (getOptional l)
+
+||| Retrieve the focus value of an optional, or allow its type to change if there
+||| is no focus.
+public export
+matching : Prism s t a b -> s -> Either t a
+matching = snd . getPrism
diff --git a/src/Control/Lens/Prism.idr b/src/Control/Lens/Prism.idr
index 040a1b3..7228e88 100644
--- a/src/Control/Lens/Prism.idr
+++ b/src/Control/Lens/Prism.idr
@@ -81,3 +81,28 @@ nearly x p = prism' (const x) (guard . p)
 public export
 only : Eq a => a -> Prism' a ()
 only x = nearly x (x ==)
+
+
+||| Create a prism that operates on `Either` values from two other prisms.
+|||
+||| This can be seen as dual to `alongside`.
+public export
+without : Prism s t a b -> Prism s' t' a' b' -> Prism (Either s s') (Either t t') (Either a a') (Either b b')
+without l l' =
+  let (inj1, prj1) = getPrism l
+      (inj2, prj2) = getPrism l'
+  in prism (bimap inj1 inj2) (either (bimap Left Left . prj1) (bimap Right Right . prj2))
+
+||| Lift a prism through a `Traversable` container.
+|||
+||| The constructed prism will only match if all of the inner elements of the
+||| `Traversable` container match.
+public export
+below : Traversable f => Prism' s a -> Prism' (f s) (f a)
+below l = withPrism l $ \inj,prj =>
+  prism (map inj) $ \s => mapFst (const s) (traverse prj s)
+
+||| Lift a prism through part of a product type.
+public export
+aside : Prism s t a b -> Prism (e, s) (e, t) (e, a) (e, b)
+aside l = withPrism l $ \inj,prj => prism (map inj) $ \(e,s) => bimap (e,) (e,) (prj s)
diff --git a/src/Control/Lens/Traversal.idr b/src/Control/Lens/Traversal.idr
index e11d9fb..1715ca7 100644
--- a/src/Control/Lens/Traversal.idr
+++ b/src/Control/Lens/Traversal.idr
@@ -5,8 +5,10 @@ import Data.Zippable
 import Data.Profunctor
 import Data.Profunctor.Traversing
 import Control.Applicative.Backwards
+import Control.Applicative.Indexing
 import Control.Lens.Optic
 import Control.Lens.Optional
+import Control.Lens.Prism
 
 %default total
 
@@ -47,6 +49,12 @@ public export
 traversed : Traversable t => Traversal (t a) (t b) a b
 traversed @{_} @{MkIsTraversal _} = traverse'
 
+||| Contstruct a traversal over a `Bitraversable` container with matching types.
+public export
+both : Bitraversable t => Traversal (t a a) (t b b) a b
+both @{_} @{MkIsTraversal _} = wander (\f => bitraverse f f)
+
+
 ||| Reverse the order of a traversal's focuses.
 public export
 backwards : Traversal s t a b -> Traversal s t a b
@@ -73,8 +81,8 @@ sequenceOf : Applicative f => Traversal s t (f a) a -> s -> f t
 sequenceOf l = traverseOf l id
 
 ||| Fold across the focuses of a traversal from the leftmost focus, providing a
-||| replacement value for each focus, and return the final accumulator along
-||| with the new structure.
+||| replacement value for each, and return the final accumulator along with the
+||| new structure.
 public export
 mapAccumLOf : Traversal s t a b -> (acc -> a -> (acc, b)) -> acc -> s -> (acc, t)
 mapAccumLOf l f z =
@@ -82,8 +90,8 @@ mapAccumLOf l f z =
   in runState z . traverseOf l g
 
 ||| Fold across the focuses of a traversal from the rightmost focus, providing a
-||| replacement value for each focus, and return the final accumulator along
-||| with the new structure.
+||| replacement value for each, and return the final accumulator along with the
+||| new structure.
 public export
 mapAccumROf : Traversal s t a b -> (acc -> a -> (acc, b)) -> acc -> s -> (acc, t)
 mapAccumROf l f z =
@@ -109,3 +117,61 @@ scanr1Of l f =
       step Nothing  x = (Just x, x)
       step (Just s) x = let r = f s x in (Just r, r)
   in snd . mapAccumROf l step Nothing
+
+
+||| Try to map over a traversal, failing if the traversal has no focuses.
+public export
+failover : Alternative f => Traversal s t a b -> (a -> b) -> s -> f t
+failover l f x =
+  let _ = Bool.Monoid.Any
+      (b, y) = traverseOf l ((True,) . f) x
+  in  guard b $> y
+
+
+||| Construct an optional that focuses on the first value of a traversal.
+|||
+||| For the fold version of this, see `pre`.
+public export
+singular : Traversal' s a -> Optional' s a
+singular l = optional' (runForget $ l (MkForget Just)) set
+  where
+    set : s -> a -> s
+    set str x = evalState True $ traverseOf l
+      (\y => if !get then put False $> x else pure y) str
+
+||| Filter the focuses of a traversal by a predicate on their ordinal positions.
+public export
+elementsOf : Traversal s t a a -> (Nat -> Bool) -> Traversal s t a a
+elementsOf l p @{MkIsTraversal _} = wander func
+  where
+    func : Applicative f => (a -> f a) -> s -> f t
+    func fn = indexing {f} (traverseOf l) $
+      \i,x => if p i then fn x else pure {f} x
+
+||| Traverse over the elements of a `Traversable` container that satisfy a
+||| predicate.
+public export
+elements : Traversable t => (Nat -> Bool) -> Traversal' (t a) a
+elements = elementsOf traversed
+
+||| Construct an optional that focuses on the nth element of a traversal.
+public export
+elementOf : Traversal' s a -> Nat -> Optional' s a
+elementOf l n = singular $ elementsOf l (n ==)
+
+||| Construct an optional that focuses on the nth element of a `Traversable`
+||| container.
+public export
+element : Traversable t => Nat -> Optional' (t a) a
+element = elementOf traversed
+
+
+||| Limit a traversal to its first n focuses.
+public export
+taking : Nat -> Traversal s t a a -> Traversal s t a a
+taking n l = elementsOf l (< n)
+
+||| Remove the first n focuses from a traversal.
+public export
+dropping : Nat -> Traversal s t a a -> Traversal s t a a
+dropping i l = elementsOf l (>= i)