add splitAt

8 files changed, 40 insertions, 0 deletions

diff --git a/blessings.cabal b/blessings.cabal
index 09b776c..d824304 100644
--- a/blessings.cabal
+++ b/blessings.cabal
@@ -9,6 +9,7 @@ library
   build-depends:
     base,
     bytestring,
+    extra,
     text
   exposed-modules:
     Blessings,
diff --git a/src/Blessings.hs b/src/Blessings.hs
index 8b55e24..59fc4c8 100644
--- a/src/Blessings.hs
+++ b/src/Blessings.hs
@@ -13,6 +13,7 @@ import Control.Applicative
 import Data.Ix (inRange)
 import Data.List (genericDrop)
 import Data.String
+import Data.Tuple.Extra (both, first, second)
 import Data.Word (Word8)
 import qualified Prelude
 import Prelude hiding (drop, length, take)
@@ -393,6 +394,22 @@ instance (Eq a, Blessable a) => Blessable (Blessings a) where
       Empty ->
           Empty
 
+  splitAt n = \case
+      Append t1 t2 ->
+        both normalize $
+          let
+            nt1 = Bless.length t1
+          in
+            if n <= nt1
+              then second (<>t2) $ Bless.splitAt n t1
+              else first (t1<>) $ Bless.splitAt (n - nt1) t2
+      Plain s ->
+        both (normalize . Plain) $ Bless.splitAt n s
+      SGR pm t ->
+        both (normalize . SGR pm) $ Bless.splitAt n t
+      Empty ->
+        (Empty, Empty)
+
   intercalate i = \case
       [] -> mempty
       [t] -> t
diff --git a/src/Blessings/ByteString.hs b/src/Blessings/ByteString.hs
index d914818..42139fa 100644
--- a/src/Blessings/ByteString.hs
+++ b/src/Blessings/ByteString.hs
@@ -13,5 +13,6 @@ instance Blessable B.ByteString where
   length = B.length
   drop = B.drop
   take = B.take
+  splitAt = B.splitAt
   intercalate = B.intercalate
   fromWord8 = B.pack . show
diff --git a/src/Blessings/ByteString/Lazy.hs b/src/Blessings/ByteString/Lazy.hs
index a32f29b..c0f521c 100644
--- a/src/Blessings/ByteString/Lazy.hs
+++ b/src/Blessings/ByteString/Lazy.hs
@@ -13,5 +13,6 @@ instance Blessable L.ByteString where
   length = fromIntegral . L.length
   drop = L.drop . fromIntegral
   take = L.take . fromIntegral
+  splitAt = L.splitAt . fromIntegral
   intercalate = L.intercalate
   fromWord8 = L.pack . show
diff --git a/src/Blessings/Internal.hs b/src/Blessings/Internal.hs
index 0ed5556..38c9069 100644
--- a/src/Blessings/Internal.hs
+++ b/src/Blessings/Internal.hs
@@ -8,5 +8,6 @@ class (IsString a, Monoid a) => Blessable a where
   length :: a -> Int
   drop :: Int -> a -> a
   take :: Int -> a -> a
+  splitAt :: Int -> a -> (a, a)
   intercalate :: a -> [a] -> a
   fromWord8 :: Word8 -> a
diff --git a/src/Blessings/String.hs b/src/Blessings/String.hs
index 005cd7b..bc3d8cc 100644
--- a/src/Blessings/String.hs
+++ b/src/Blessings/String.hs
@@ -14,5 +14,6 @@ instance Blessable String where
   length = L.length
   drop = L.drop
   take = L.take
+  splitAt = L.splitAt
   intercalate = L.intercalate
   fromWord8 = show
diff --git a/src/Blessings/Text.hs b/src/Blessings/Text.hs
index 1f82c22..236b5d0 100644
--- a/src/Blessings/Text.hs
+++ b/src/Blessings/Text.hs
@@ -14,5 +14,6 @@ instance Blessable Text where
   length = T.length
   drop = T.drop
   take = T.take
+  splitAt = T.splitAt
   intercalate = T.intercalate
   fromWord8 = T.pack . show
diff --git a/test/Spec.hs b/test/Spec.hs
index 6cf4d2f..24a17e2 100644
--- a/test/Spec.hs
+++ b/test/Spec.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE MultiWayIf #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 import Blessings.Internal as B
@@ -43,6 +44,22 @@ main =
         property $ \(x :: Blessings String) ->
           normalize (B.take 1 x <> B.drop 1 x) == normalize x
 
+      it "uncurry (<>) (splitAt i x) == x" $
+        property $ \(i :: Int, x :: Blessings String) ->
+          unsafeTimeout 100000 $
+            normalize (uncurry (<>) (B.splitAt i x)) == normalize x
+
+      it "splitAt produces pairs with elements of proper length" $
+        property $ \(i :: Int, x :: Blessings String) ->
+          unsafeTimeout 100000 $
+            let
+              (l, r) = B.splitAt i x
+              n = B.length x
+            in
+            if | i <= 0    -> B.length l == 0 && B.length r == n
+               | i <= n    -> B.length l == i && B.length r == n - i
+               | otherwise -> B.length l == n && B.length r == 0
+
       let infx = mconcat (repeat (Plain "x" :: Blessings String))
 
       it "can take from infinite structure" $