{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleContexts #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Bits.Lens
-- Copyright   :  (C) 2012-16 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  LiberalTypeSynonyms
--
----------------------------------------------------------------------------
module Data.Bits.Lens
  ( (.|.~), (.&.~), (<.|.~), (<.&.~), (<<.|.~), (<<.&.~)
  , (.|.=), (.&.=), (<.|.=), (<.&.=), (<<.|.=), (<<.&.=)
  , bitAt
  , bits
  , byteAt
  , bytewise
  ) where

import Prelude ()

import Control.Lens
import Control.Lens.Internal.Prelude
import Control.Monad.State
import Data.Bits
import Data.Word

-- $setup
-- >>> :set -XNoOverloadedStrings
-- >>> import Control.Lens
-- >>> import Control.Monad.State
-- >>> import Data.Word

infixr 4 .|.~, .&.~, <.|.~, <.&.~, <<.|.~, <<.&.~
infix 4 .|.=, .&.=, <.|.=, <.&.=, <<.|.=, <<.&.=

-- | Bitwise '.|.' the target(s) of a t'Lens' or t'Setter'.
--
-- >>> _2 .|.~ 6 $ ("hello",3)
-- ("hello",7)
--
-- @
-- ('.|.~') :: 'Bits' a             => t'Setter' s t a a    -> a -> s -> t
-- ('.|.~') :: 'Bits' a             => t'Iso' s t a a       -> a -> s -> t
-- ('.|.~') :: 'Bits' a             => t'Lens' s t a a      -> a -> s -> t
-- ('.|.~') :: ('Data.Monoid.Monoid' a, 'Bits' a) => t'Traversal' s t a a -> a -> s -> t
-- @
(.|.~):: Bits a => ASetter s t a a -> a -> s -> t
ASetter s t a a
l .|.~ :: forall a s t. Bits a => ASetter s t a a -> a -> s -> t
.|.~ a
n = forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter s t a a
l (forall a. Bits a => a -> a -> a
.|. a
n)
{-# INLINE (.|.~) #-}

-- | Bitwise '.&.' the target(s) of a t'Lens' or t'Setter'.
--
-- >>> _2 .&.~ 7 $ ("hello",254)
-- ("hello",6)
--
-- @
-- ('.&.~') :: 'Bits' a             => t'Setter' s t a a    -> a -> s -> t
-- ('.&.~') :: 'Bits' a             => t'Iso' s t a a       -> a -> s -> t
-- ('.&.~') :: 'Bits' a             => t'Lens' s t a a      -> a -> s -> t
-- ('.&.~') :: ('Data.Monoid.Monoid' a, 'Bits' a) => t'Traversal' s t a a -> a -> s -> t
-- @
(.&.~) :: Bits a => ASetter s t a a -> a -> s -> t
ASetter s t a a
l .&.~ :: forall a s t. Bits a => ASetter s t a a -> a -> s -> t
.&.~ a
n = forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter s t a a
l (forall a. Bits a => a -> a -> a
.&. a
n)
{-# INLINE (.&.~) #-}

-- | Modify the target(s) of a 'Lens'', 'Setter'' or 'Traversal'' by computing its bitwise '.&.' with another value.
--
-- >>> execState (do _1 .&.= 15; _2 .&.= 3) (7,7)
-- (7,3)
--
-- @
-- ('.&.=') :: ('MonadState' s m, 'Bits' a) => 'Setter'' s a    -> a -> m ()
-- ('.&.=') :: ('MonadState' s m, 'Bits' a) => 'Iso'' s a       -> a -> m ()
-- ('.&.=') :: ('MonadState' s m, 'Bits' a) => 'Lens'' s a      -> a -> m ()
-- ('.&.=') :: ('MonadState' s m, 'Bits' a) => 'Traversal'' s a -> a -> m ()
-- @
(.&.=):: (MonadState s m, Bits a) => ASetter' s a -> a -> m ()
ASetter' s a
l .&.= :: forall s (m :: * -> *) a.
(MonadState s m, Bits a) =>
ASetter' s a -> a -> m ()
.&.= a
a = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (ASetter' s a
l forall a s t. Bits a => ASetter s t a a -> a -> s -> t
.&.~ a
a)
{-# INLINE (.&.=) #-}

-- | Modify the target(s) of a 'Lens'', t'Setter' or t'Traversal' by computing its bitwise '.|.' with another value.
--
-- >>> execState (do _1 .|.= 15; _2 .|.= 3) (7,7)
-- (15,7)
--
-- @
-- ('.|.=') :: ('MonadState' s m, 'Bits' a) => 'Setter'' s a    -> a -> m ()
-- ('.|.=') :: ('MonadState' s m, 'Bits' a) => 'Iso'' s a       -> a -> m ()
-- ('.|.=') :: ('MonadState' s m, 'Bits' a) => 'Lens'' s a      -> a -> m ()
-- ('.|.=') :: ('MonadState' s m, 'Bits' a) => 'Traversal'' s a -> a -> m ()
-- @
(.|.=) :: (MonadState s m, Bits a) => ASetter' s a -> a -> m ()
ASetter' s a
l .|.= :: forall s (m :: * -> *) a.
(MonadState s m, Bits a) =>
ASetter' s a -> a -> m ()
.|.= a
a = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (ASetter' s a
l forall a s t. Bits a => ASetter s t a a -> a -> s -> t
.|.~ a
a)
{-# INLINE (.|.=) #-}

-- | Bitwise '.|.' the target(s) of a t'Lens' (or t'Traversal'), returning the result
-- (or a monoidal summary of all of the results).
--
-- >>> _2 <.|.~ 6 $ ("hello",3)
-- (7,("hello",7))
--
-- @
-- ('<.|.~') :: 'Bits' a             => t'Iso' s t a a       -> a -> s -> (a, t)
-- ('<.|.~') :: 'Bits' a             => t'Lens' s t a a      -> a -> s -> (a, t)
-- ('<.|.~') :: ('Bits' a, 'Data.Monoid.Monoid' a) => t'Traversal' s t a a -> a -> s -> (a, t)
-- @
(<.|.~):: Bits a => LensLike ((,) a) s t a a -> a -> s -> (a, t)
LensLike ((,) a) s t a a
l <.|.~ :: forall a s t.
Bits a =>
LensLike ((,) a) s t a a -> a -> s -> (a, t)
<.|.~ a
n = LensLike ((,) a) s t a a
l forall b s t a. LensLike ((,) b) s t a b -> (a -> b) -> s -> (b, t)
<%~ (forall a. Bits a => a -> a -> a
.|. a
n)
{-# INLINE (<.|.~) #-}

-- | Bitwise '.&.' the target(s) of a t'Lens' or t'Traversal', returning the result
-- (or a monoidal summary of all of the results).
--
-- >>> _2 <.&.~ 7 $ ("hello",254)
-- (6,("hello",6))
--
-- @
-- ('<.&.~') :: 'Bits' a             => t'Iso'       s t a a -> a -> s -> (a, t)
-- ('<.&.~') :: 'Bits' a             => t'Lens'      s t a a -> a -> s -> (a, t)
-- ('<.&.~') :: ('Bits' a, 'Data.Monoid.Monoid' a) => t'Traversal' s t a a -> a -> s -> (a, t)
-- @
(<.&.~) :: Bits a => LensLike ((,) a) s t a a -> a -> s -> (a, t)
LensLike ((,) a) s t a a
l <.&.~ :: forall a s t.
Bits a =>
LensLike ((,) a) s t a a -> a -> s -> (a, t)
<.&.~ a
n = LensLike ((,) a) s t a a
l forall b s t a. LensLike ((,) b) s t a b -> (a -> b) -> s -> (b, t)
<%~ (forall a. Bits a => a -> a -> a
.&. a
n)
{-# INLINE (<.&.~) #-}

-- | Modify the target(s) of a 'Lens'' (or 'Traversal'') by computing its bitwise '.&.' with another value,
-- returning the result (or a monoidal summary of all of the results traversed).
--
-- >>> runState (_1 <.&.= 15) (31,0)
-- (15,(15,0))
--
-- @
-- ('<.&.=') :: ('MonadState' s m, 'Bits' a)           => 'Lens'' s a      -> a -> m a
-- ('<.&.=') :: ('MonadState' s m, 'Bits' a, 'Data.Monoid.Monoid' a) => 'Traversal'' s a -> a -> m a
-- @
(<.&.=):: (MonadState s m, Bits a) => LensLike' ((,)a) s a -> a -> m a
LensLike' ((,) a) s a
l <.&.= :: forall s (m :: * -> *) a.
(MonadState s m, Bits a) =>
LensLike' ((,) a) s a -> a -> m a
<.&.= a
b = LensLike' ((,) a) s a
l forall s (m :: * -> *) b a.
MonadState s m =>
LensLike ((,) b) s s a b -> (a -> b) -> m b
<%= (forall a. Bits a => a -> a -> a
.&. a
b)
{-# INLINE (<.&.=) #-}

-- | Modify the target(s) of a 'Lens'', (or t'Traversal') by computing its bitwise '.|.' with another value,
-- returning the result (or a monoidal summary of all of the results traversed).
--
-- >>> runState (_1 <.|.= 7) (28,0)
-- (31,(31,0))
--
-- @
-- ('<.|.=') :: ('MonadState' s m, 'Bits' a)           => 'Lens'' s a      -> a -> m a
-- ('<.|.=') :: ('MonadState' s m, 'Bits' a, 'Data.Monoid.Monoid' a) => 'Traversal'' s a -> a -> m a
-- @
(<.|.=) :: (MonadState s m, Bits a) => LensLike' ((,)a) s a -> a -> m a
LensLike' ((,) a) s a
l <.|.= :: forall s (m :: * -> *) a.
(MonadState s m, Bits a) =>
LensLike' ((,) a) s a -> a -> m a
<.|.= a
b = LensLike' ((,) a) s a
l forall s (m :: * -> *) b a.
MonadState s m =>
LensLike ((,) b) s s a b -> (a -> b) -> m b
<%= (forall a. Bits a => a -> a -> a
.|. a
b)
{-# INLINE (<.|.=) #-}

-- | Bitwise '.&.' the target(s) of a t'Lens' or t'Traversal', and return the
-- original value, or a monoidal summary of the original values.
--
-- When you do not need the old value, ('.&.~') is more flexible.
--
-- >>> _2 <<.&.~ 7 $ ("hello", 254)
-- (254,("hello",6))
--
-- @
-- ('<<.&.~') ::  'Bits' a            => t'Iso' s t a a       -> a -> s -> (a, t)
-- ('<<.&.~') ::  'Bits' a            => t'Lens' s t a a      -> a -> s -> (a, t)
-- ('<<.&.~') :: ('Bits' a, 'Data.Monoid.Monoid' a) => t'Traversal' s t a a -> a -> s -> (a, t)
-- @
(<<.&.~) :: Bits a => Optical' (->) q ((,)a) s a -> a -> q s (a, s)
Optical' (->) q ((,) a) s a
l <<.&.~ :: forall a (q :: * -> * -> *) s.
Bits a =>
Optical' (->) q ((,) a) s a -> a -> q s (a, s)
<<.&.~ a
b = Optical' (->) q ((,) a) s a
l forall a b. (a -> b) -> a -> b
$ \a
a -> (a
a, a
a forall a. Bits a => a -> a -> a
.&. a
b)
{-# INLINE (<<.&.~) #-}

-- | Bitwise '.|.' the target(s) of a t'Lens' or t'Traversal', and return the
-- original value, or a monoidal summary of the original values.
--
-- When you do not need the old value, ('.|.~') is more flexible.
--
-- >>> _2 <<.|.~ 6 $ ("hello", 3)
-- (3,("hello",7))
--
-- @
-- ('<<.|.~') ::  'Bits' a            => t'Iso' s t a a       -> a -> s -> (a, t)
-- ('<<.|.~') ::  'Bits' a            => t'Lens' s t a a      -> a -> s -> (a, t)
-- ('<<.|.~') :: ('Bits' a, 'Data.Monoid.Monoid' a) => t'Traversal' s t a a -> a -> s -> (a, t)
-- @
(<<.|.~) :: Bits a => Optical' (->) q ((,)a) s a -> a -> q s (a, s)
Optical' (->) q ((,) a) s a
l <<.|.~ :: forall a (q :: * -> * -> *) s.
Bits a =>
Optical' (->) q ((,) a) s a -> a -> q s (a, s)
<<.|.~ a
b = Optical' (->) q ((,) a) s a
l forall a b. (a -> b) -> a -> b
$ \a
a -> (a
a, a
a forall a. Bits a => a -> a -> a
.|. a
b)
{-# INLINE (<<.|.~) #-}

-- | Modify the target(s) of a 'Lens'', (or 'Traversal'') by computing its
-- bitwise '.&.' with another value, returning the original value (or a
-- monoidal summary of all the original values).
--
-- When you do not need the old value, ('.&.=') is more flexible.
--
-- >>> runState (_1 <<.&.= 15) (31,0)
-- (31,(15,0))
--
-- @
-- ('<<.&.=') :: ('MonadState' s m, 'Bits' a)           => 'Lens'' s a      -> a -> m a
-- ('<<.&.=') :: ('MonadState' s m, 'Bits' a, 'Data.Monoid.Monoid' a) => 'Traversal'' s a -> a -> m a
-- @
(<<.&.=) :: (MonadState s m, Bits a) => LensLike' ((,) a) s a -> a -> m a
LensLike' ((,) a) s a
l <<.&.= :: forall s (m :: * -> *) a.
(MonadState s m, Bits a) =>
LensLike' ((,) a) s a -> a -> m a
<<.&.= a
b = LensLike' ((,) a) s a
l forall {k} s (m :: * -> *) (p :: k -> * -> *) r (a :: k) b.
MonadState s m =>
Over p ((,) r) s s a b -> p a (r, b) -> m r
%%= \a
a -> (a
a, a
a forall a. Bits a => a -> a -> a
.&. a
b)
{-# INLINE (<<.&.=) #-}

-- | Modify the target(s) of a 'Lens'', (or 'Traversal'') by computing its
-- bitwise '.|.' with another value, returning the original value (or a
-- monoidal summary of all the original values).
--
-- When you do not need the old value, ('.|.=') is more flexible.
--
-- >>> runState (_1 <<.|.= 7) (28,0)
-- (28,(31,0))
--
-- @
-- ('<<.|.=') :: ('MonadState' s m, 'Bits' a)           => 'Lens'' s a      -> a -> m a
-- ('<<.|.=') :: ('MonadState' s m, 'Bits' a, 'Data.Monoid.Monoid' a) => 'Traversal'' s a -> a -> m a
-- @
(<<.|.=) :: (MonadState s m, Bits a) => LensLike' ((,) a) s a -> a -> m a
LensLike' ((,) a) s a
l <<.|.= :: forall s (m :: * -> *) a.
(MonadState s m, Bits a) =>
LensLike' ((,) a) s a -> a -> m a
<<.|.= a
b = LensLike' ((,) a) s a
l forall {k} s (m :: * -> *) (p :: k -> * -> *) r (a :: k) b.
MonadState s m =>
Over p ((,) r) s s a b -> p a (r, b) -> m r
%%= \a
a -> (a
a, a
a forall a. Bits a => a -> a -> a
.|. a
b)
{-# INLINE (<<.|.=) #-}

-- | This t'Lens' can be used to access the value of the nth bit in a number.
--
-- @'bitAt' n@ is only a legal t'Lens' into @b@ if @0 '<=' n '<' 'bitSize' ('undefined' :: b)@.
--
-- >>> 16^.bitAt 4
-- True
--
-- >>> 15^.bitAt 4
-- False
--
-- >>> 15 & bitAt 4 .~ True
-- 31
--
-- >>> 16 & bitAt 4 .~ False
-- 0
bitAt :: Bits b => Int -> IndexedLens' Int b Bool
bitAt :: forall b. Bits b => Int -> IndexedLens' Int b Bool
bitAt Int
n p Bool (f Bool)
f b
b = forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p Bool (f Bool)
f Int
n (forall a. Bits a => a -> Int -> Bool
testBit b
b Int
n) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Bool
x -> if Bool
x then forall a. Bits a => a -> Int -> a
setBit b
b Int
n else forall a. Bits a => a -> Int -> a
clearBit b
b Int
n
{-# INLINE bitAt #-}

-- | Get the nth byte, counting from the low end.
--
-- @'byteAt' n@ is a legal t'Lens' into @b@ iff @0 '<=' n '<' 'div' ('bitSize' ('undefined' :: b)) 8@
--
-- >>> (0xff00 :: Word16)^.byteAt 0
-- 0
--
-- >>> (0xff00 :: Word16)^.byteAt 1
-- 255
--
-- >>> byteAt 1 .~ 0 $ 0xff00 :: Word16
-- 0
--
-- >>> byteAt 0 .~ 0xff $ 0 :: Word16
-- 255
byteAt :: (Integral b, Bits b) => Int -> IndexedLens' Int b Word8
byteAt :: forall b. (Integral b, Bits b) => Int -> IndexedLens' Int b Word8
byteAt Int
i p Word8 (f Word8)
f b
b = forall {a}. Integral a => a -> b
back forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p Word8 (f Word8)
f Int
i (b -> Word8
forward b
b) where
  back :: a -> b
back a
w8 = (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w8 forall a. Bits a => a -> Int -> a
`shiftL` (Int
i forall a. Num a => a -> a -> a
* Int
8))
    forall a. Bits a => a -> a -> a
.|. (forall a. Bits a => a -> a
complement (b
255 forall a. Bits a => a -> Int -> a
`shiftL` (Int
i forall a. Num a => a -> a -> a
* Int
8)) forall a. Bits a => a -> a -> a
.&. b
b)
  forward :: b -> Word8
forward = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bits a => a -> a -> a
(.&.) b
0xff forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Bits a => a -> Int -> a
shiftR (Int
i forall a. Num a => a -> a -> a
* Int
8)


-- | Traverse over all bits in a numeric type.
--
-- The bit position is available as the index.
--
-- >>> toListOf bits (5 :: Word8)
-- [True,False,True,False,False,False,False,False]
--
-- If you supply this an 'Integer', the result will be an infinite t'Traversal', which
-- can be productively consumed, but not reassembled.
bits :: (Num b, Bits b) => IndexedTraversal' Int b Bool
bits :: forall b. (Num b, Bits b) => IndexedTraversal' Int b Bool
bits p Bool (f Bool)
f b
b = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall {a}. Bits a => (Int, Bool) -> a -> a
step b
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Int -> f (Int, Bool)
g [Int]
bs where
  g :: Int -> f (Int, Bool)
g Int
n      = (,) Int
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p Bool (f Bool)
f Int
n (forall a. Bits a => a -> Int -> Bool
testBit b
b Int
n)
  bs :: [Int]
bs       = forall a. (a -> Bool) -> [a] -> [a]
takeWhile Int -> Bool
hasBit [Int
0..]
  hasBit :: Int -> Bool
hasBit Int
n = forall a. Bits a => a -> Int -> a
complementBit b
b Int
n forall a. Eq a => a -> a -> Bool
/= b
b -- test to make sure that complementing this bit actually changes the value
  step :: (Int, Bool) -> a -> a
step (Int
n,Bool
True) a
r = forall a. Bits a => a -> Int -> a
setBit a
r Int
n
  step (Int, Bool)
_        a
r = a
r
{-# INLINE bits #-}

-- | Traverse over all the bytes in an integral type, from the low end.
--
-- The byte position is available as the index.
--
-- >>> toListOf bytewise (1312301580 :: Word32)
-- [12,34,56,78]
--
-- If you supply this an 'Integer', the result will be an infinite t'Traversal',
-- which can be productively consumed, but not reassembled.
--
-- Why isn't this function called @bytes@ to match 'bits'? Alas, there
-- is already a function by that name in "Data.ByteString.Lens".
bytewise :: (Integral b, Bits b) => IndexedTraversal' Int b Word8
bytewise :: forall b. (Integral b, Bits b) => IndexedTraversal' Int b Word8
bytewise p Word8 (f Word8)
f b
b = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall {a} {a}. (Bits a, Integral a, Num a) => (Int, a) -> a -> a
step b
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Int -> f (Int, Word8)
g [Int]
bs where
  g :: Int -> f (Int, Word8)
g Int
n = (,) Int
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p Word8 (f Word8)
f Int
n (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ b
b forall a. Bits a => a -> Int -> a
`shiftR` (Int
nforall a. Num a => a -> a -> a
*Int
8))
  bs :: [Int]
bs = forall a. (a -> Bool) -> [a] -> [a]
takeWhile Int -> Bool
hasByte [Int
0..]
  hasByte :: Int -> Bool
hasByte Int
n = forall a. Bits a => a -> Int -> a
complementBit b
b (Int
nforall a. Num a => a -> a -> a
*Int
8) forall a. Eq a => a -> a -> Bool
/= b
b
  step :: (Int, a) -> a -> a
step (Int
n,a
x) a
r = a
r forall a. Bits a => a -> a -> a
.|. (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x forall a. Bits a => a -> Int -> a
`shiftL` (Int
nforall a. Num a => a -> a -> a
*Int
8))
{-# INLINE bytewise #-}