more efficient traversal

Author: b-mehta

src/Control/Optics/Linear/Internal.hs

@@ -189,15 +189,15 @@ traversal :: (s ->. Batch a b t) -> Traversal a b s t
 traversal h = Optical (\k -> dimap h fuse (traverse' k))
 
 traverse' :: (Strong Either Void arr, Monoidal (,) () arr) => a `arr` b -> Batch a c t `arr` Batch b c t
-traverse' k = dimap out inn (second (traverse' k *** k))
+traverse' k = dimap fromBatch toBatch (second (k *** traverse' k))
 
-out :: Batch a b t ->. Either t (Batch a b (b ->. t), a)
-out (P t) = Left t
-out (l :*: x) = Right (l,x)
+fromBatch :: Batch a b t ->. Either t (a, Batch a b (b ->. t))
+fromBatch (Done t) = Left t
+fromBatch (More l x) = Right (l, x)
 
-inn :: Either t (Batch a b (b ->. t), a) ->. Batch a b t
-inn (Left t) = P t
-inn (Right (l,x)) = l :*: x
+toBatch :: Either t (a, Batch a b (b ->. t)) ->. Batch a b t
+toBatch (Left t) = Done t
+toBatch (Right (l, x)) = More l x
 
 traversed :: Traversable t => Traversal a b (t a) (t b)
 traversed = traversal (traverse batch)

src/Data/Functor/Linear/Internal/Traversable.hs

@@ -80,31 +80,28 @@ instance Control.Applicative (StateR s) where
     where go :: (a, (a ->. b, s)) ->. (b, s)
           go (a, (h, s'')) = (h a, s'')
 
-data Batch a b c = P c | Batch a b (b ->. c) :*: a
+data Batch a b c = Done c | More a (Batch a b (b ->. c))
   deriving (Data.Functor, Data.Applicative) via Control.Data (Batch a b)
+
 instance Control.Functor (Batch a b) where
-  fmap f (P c) = P (f c)
-  fmap f (u :*: a) = Control.fmap (f.) u :*: a
+  fmap f (Done c)   = Done (f c)
+  fmap f (More x l) = More x ((f.) Control.<$> l)
 
 instance Control.Applicative (Batch a b) where
-  pure = P
-  P f <*> P x = P (f x)
-  (u :*: a) <*> P x = ((P $ help x) Control.<*> u) :*: a
-  u <*> (v :*: a) = (P (.) Control.<*> u Control.<*> v) :*: a
-
-help :: d ->. ((b ->. d ->. e) ->. b ->. e)
-help d bde b = bde b d
+  pure = Done
+  Done f <*> l' = Control.fmap f l'
+  More x l <*> l' = More x (flip Control.<$> l Control.<*> l')
 
 batch :: a ->. Batch a b b
-batch x = P id :*: x
+batch x = More x (Done id)
 
 runWith :: Control.Applicative f => (a ->. f b) -> Batch a b c ->. f c
-runWith _ (P x) = Control.pure x
-runWith f (u :*: x) = runWith f u Control.<*> f x
+runWith _ (Done c) = Control.pure c
+runWith f (More x l) = runWith f l Control.<*> f x
 
 fuse :: Batch b b t ->. t
-fuse (P i) = i
-fuse (u :*: x) = fuse u x
+fuse (Done i) = i
+fuse (More x l) = fuse l x
 
 ------------------------
 -- Standard instances --

src/Prelude/Linear.hs

@@ -83,10 +83,6 @@ maybe _ f (Just y) = f y
 forget :: (a ->. b) ->. a -> b
 forget f x = f x
 
--- | Replacement for the flip function with generalized multiplicities.
-flip :: (a -->.(p) b -->.(q) c) -->.(r) b -->.(q) a -->.(p) c
-flip f b a = f a b
-
 -- | Linearly typed replacement for the standard '(Prelude.<*)' function.
 (<*) :: (Data.Applicative f, Consumable b) => f a ->. f b ->. f a
 fa <* fb = Data.fmap (flip lseq) fa Data.<*> fb

src/Prelude/Linear/Internal/Simple.hs

@@ -56,3 +56,7 @@ foldr :: (a ->. b ->. b) -> b ->. [a] ->. b
 foldr f z = \case
   [] -> z
   x:xs -> f x (foldr f z xs)
+
+-- | Replacement for the flip function with generalized multiplicities.
+flip :: (a -->.(p) b -->.(q) c) -->.(r) b -->.(q) a -->.(p) c
+flip f b a = f a b