{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE RankNTypes #-}

-- | Reimplementations of stuff otherwise found in other libraries
--
-- "A little copying is better than a little dependency"
module Nonlinear.Internal
  ( imap,
    Lens',
    set,
    view,
    lens,
    ASetter',
  )
where

import Control.Applicative (Const (Const, getConst))
import Control.Monad (ap)
import Data.Functor.Identity (Identity (runIdentity))
import Data.Traversable (for)

{-# INLINE imap #-}
imap :: Traversable t => (Int -> a -> b) -> (t a -> t b)
imap :: (Int -> a -> b) -> t a -> t b
imap Int -> a -> b
f t a
t = Tally (t b) -> t b
forall a. Tally a -> a
evalTally (Tally (t b) -> t b)
-> ((a -> Tally b) -> Tally (t b)) -> (a -> Tally b) -> t b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t a -> (a -> Tally b) -> Tally (t b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for t a
t ((a -> Tally b) -> t b) -> (a -> Tally b) -> t b
forall a b. (a -> b) -> a -> b
$ \a
a -> (Int -> a -> b) -> a -> Int -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> a -> b
f a
a (Int -> b) -> Tally Int -> Tally b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tally Int
click

-- | Equivalent to State Int
newtype Tally a = Tally {Tally a -> Int -> (Int, a)
unTally :: Int -> (Int, a)}
  deriving (a -> Tally b -> Tally a
(a -> b) -> Tally a -> Tally b
(forall a b. (a -> b) -> Tally a -> Tally b)
-> (forall a b. a -> Tally b -> Tally a) -> Functor Tally
forall a b. a -> Tally b -> Tally a
forall a b. (a -> b) -> Tally a -> Tally b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Tally b -> Tally a
$c<$ :: forall a b. a -> Tally b -> Tally a
fmap :: (a -> b) -> Tally a -> Tally b
$cfmap :: forall a b. (a -> b) -> Tally a -> Tally b
Functor)

instance Applicative Tally where
  {-# INLINE pure #-}
  pure :: a -> Tally a
pure a
a = (Int -> (Int, a)) -> Tally a
forall a. (Int -> (Int, a)) -> Tally a
Tally ((Int -> (Int, a)) -> Tally a) -> (Int -> (Int, a)) -> Tally a
forall a b. (a -> b) -> a -> b
$ \Int
i -> (Int
i, a
a)
  {-# INLINE (<*>) #-}
  <*> :: Tally (a -> b) -> Tally a -> Tally b
(<*>) = Tally (a -> b) -> Tally a -> Tally b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad Tally where
  {-# INLINE (>>=) #-}
  Tally Int -> (Int, a)
sa >>= :: Tally a -> (a -> Tally b) -> Tally b
>>= a -> Tally b
asb = (Int -> (Int, b)) -> Tally b
forall a. (Int -> (Int, a)) -> Tally a
Tally ((Int -> (Int, b)) -> Tally b) -> (Int -> (Int, b)) -> Tally b
forall a b. (a -> b) -> a -> b
$ \Int
i -> let (Int
i', a
a) = Int -> (Int, a)
sa Int
i in Tally b -> Int -> (Int, b)
forall a. Tally a -> Int -> (Int, a)
unTally (a -> Tally b
asb a
a) Int
i'

{-# INLINE click #-}
click :: Tally Int
click :: Tally Int
click = (Int -> (Int, Int)) -> Tally Int
forall a. (Int -> (Int, a)) -> Tally a
Tally ((Int -> (Int, Int)) -> Tally Int)
-> (Int -> (Int, Int)) -> Tally Int
forall a b. (a -> b) -> a -> b
$ \Int
i -> (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
i)

{-# INLINE evalTally #-}
evalTally :: Tally a -> a
evalTally :: Tally a -> a
evalTally = (Int, a) -> a
forall a b. (a, b) -> b
snd ((Int, a) -> a) -> (Tally a -> (Int, a)) -> Tally a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Tally a -> Int -> (Int, a)) -> Int -> Tally a -> (Int, a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Tally a -> Int -> (Int, a)
forall a. Tally a -> Int -> (Int, a)
unTally Int
0

type Lens' s a = forall m. Functor m => (a -> m a) -> (s -> m s)

{-# INLINE set #-}
set :: ASetter' s a -> a -> s -> s
set :: ASetter' s a -> a -> s -> s
set ASetter' s a
l a
a s
s = Identity s -> s
forall a. Identity a -> a
runIdentity (Identity s -> s) -> Identity s -> s
forall a b. (a -> b) -> a -> b
$ ASetter' s a
l (Identity a -> a -> Identity a
forall a b. a -> b -> a
const (Identity a -> a -> Identity a) -> Identity a -> a -> Identity a
forall a b. (a -> b) -> a -> b
$ a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a) s
s

{-# INLINE view #-}
view :: Lens' s a -> s -> a
view :: Lens' s a -> s -> a
view Lens' s a
l s
s = Const a s -> a
forall a k (b :: k). Const a b -> a
getConst (Const a s -> a) -> Const a s -> a
forall a b. (a -> b) -> a -> b
$ (a -> Const a a) -> s -> Const a s
Lens' s a
l a -> Const a a
forall k a (b :: k). a -> Const a b
Const s
s

{-# INLINE lens #-}
lens :: (s -> a) -> (s -> a -> s) -> Lens' s a
lens :: (s -> a) -> (s -> a -> s) -> Lens' s a
lens s -> a
sa s -> a -> s
sab a -> m a
l s
s = s -> a -> s
sab s
s (a -> s) -> m a -> m s
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m a
l (s -> a
sa s
s)

type ASetter' s a = (a -> Identity a) -> (s -> Identity s)