Speed up readDriveLetter 10x by introducing uncons2

Author: Bodigrim

System/FilePath/Internal.hs

@@ -506,12 +506,14 @@ readDriveUNC bs = case unpack bs of
 
 {- c:\ -}
 readDriveLetter :: STRING -> Maybe (FILEPATH, FILEPATH)
-readDriveLetter bs = case unpack bs of
-  (x:c:y:xs)
-    | c == _colon && isLetter x && isPathSeparator y -> Just $ addSlash (pack [x,_colon]) (pack (y:xs))
-  (x:c:xs)
-    | c == _colon && isLetter x -> Just (pack [x,_colon], pack xs)
-  _ -> Nothing
+readDriveLetter bs = case uncons2 bs of
+  Nothing -> Nothing
+  Just (x, c, ys)
+    | isLetter x, c == _colon -> Just $ case uncons ys of
+      Just (y, _)
+        | isPathSeparator y -> addSlash (pack [x,_colon]) ys
+      _ -> (pack [x,_colon], ys)
+    | otherwise -> Nothing
 
 {- \\sharename\ -}
 readDriveShare :: STRING -> Maybe (FILEPATH, FILEPATH)
@@ -1147,6 +1149,10 @@ stripSuffix xs ys = reverse P.<$> stripPrefix (reverse xs) (reverse ys)
 unsnoc :: [a] -> Maybe ([a], a)
 unsnoc = L.foldr (\x -> Just . maybe ([], x) (first (x :))) Nothing
 
+uncons2 :: [a] -> Maybe (a, a, [a])
+uncons2 [] = Nothing
+uncons2 [_] = Nothing
+uncons2 (x : y : zs) = Just (x, y, zs)
 
 _period, _quotedbl, _backslash, _slash, _question, _U, _N, _C, _colon, _semicolon, _US, _less, _greater, _bar, _asterisk, _nul, _space, _underscore :: Char
 _period = '.'

System/OsPath/Data/ByteString/Short.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE MultiWayIf #-}
 {-# LANGUAGE NoImplicitPrelude #-}
 -- |
 -- Module      : System.OsPath.Data.ByteString.Short
@@ -81,6 +82,7 @@ module System.OsPath.Data.ByteString.Short (
     last,
     tail,
     uncons,
+    uncons2,
     head,
     init,
     unsnoc,
@@ -173,3 +175,17 @@ module System.OsPath.Data.ByteString.Short (
   ) where
 
 import Data.ByteString.Short.Internal
+import System.OsPath.Data.ByteString.Short.Internal
+
+import Prelude (Maybe(..), Ord(..), Num(..), ($), otherwise)
+import Data.Word (Word8)
+
+uncons2 :: ShortByteString -> Maybe (Word8, Word8, ShortByteString)
+uncons2 = \sbs ->
+  let l  = length sbs
+      nl = l - 2
+  in if | l <= 1 -> Nothing
+        | otherwise -> let h  = indexWord8Array (asBA sbs) 0
+                           h' = indexWord8Array (asBA sbs) 1
+                           t  = create nl $ \mba -> copyByteArray (asBA sbs) 1 mba 0 nl
+                       in Just (h, h', t)

System/OsPath/Data/ByteString/Short/Internal.hs

@@ -20,6 +20,7 @@ module System.OsPath.Data.ByteString.Short.Internal where
 
 import Control.Monad.ST
 import Control.Exception (assert, throwIO)
+import Data.Bits (Bits(..))
 import Data.ByteString.Short.Internal (ShortByteString(..), length)
 #if !MIN_VERSION_base(4,11,0)
 import Data.Semigroup
@@ -284,15 +285,21 @@ writeWord16Array (MBA# mba#) (I# i#) (W16# w#) =
       ST (\s -> case writeWord8Array# mba# (i# +# 1#) msb# s of
           s' -> (# s', () #))
 
+indexWord8Array :: BA
+                -> Int      -- ^ Word8 index
+                -> Word8
+indexWord8Array (BA# ba#) (I# i#) = W8# (indexWord8Array# ba# i#)
+
 -- | This isn't strictly Word16 array read. Instead it's two Word8 array reads
 -- to avoid endianness issues due to primops doing automatic alignment based
 -- on host platform. We expect the byte array to be LE always.
 indexWord16Array :: BA
                  -> Int      -- ^ Word8 index (not Word16)
                  -> Word16
-indexWord16Array (BA# ba#) (I# i#) =
-  case (# indexWord8Array# ba# i#, indexWord8Array# ba# (i# +# 1#) #) of
-    (# lsb#, msb# #) -> W16# (decodeWord16LE# (# lsb#, msb# #))
+indexWord16Array ba i = fromIntegral lsb .|. (fromIntegral msb `shiftL` 8)
+  where
+    lsb = indexWord8Array ba i
+    msb = indexWord8Array ba (i + 1)
 
 #if !MIN_VERSION_base(4,16,0)
 

System/OsPath/Data/ByteString/Short/Word16.hs

@@ -46,6 +46,7 @@ module System.OsPath.Data.ByteString.Short.Word16 (
     last,
     tail,
     uncons,
+    uncons2,
     head,
     init,
     unsnoc,
@@ -260,6 +261,18 @@ uncons = \(assertEven -> sbs) ->
                            t = create nl $ \mba -> copyByteArray (asBA sbs) 2 mba 0 nl
                        in Just (h, t)
 
+-- | /O(n)/ Extract first two elements and the rest of a ByteString,
+-- returning Nothing if it is shorter than two elements.
+uncons2 :: ShortByteString -> Maybe (Word16, Word16, ShortByteString)
+uncons2 = \(assertEven -> sbs) ->
+  let l  = BS.length sbs
+      nl = l - 4
+  in if | l <= 2 -> Nothing
+        | otherwise -> let h  = indexWord16Array (asBA sbs) 0
+                           h' = indexWord16Array (asBA sbs) 2
+                           t  = create nl $ \mba -> copyByteArray (asBA sbs) 4 mba 0 nl
+                       in Just (h, h', t)
+
 -- | /O(1)/ Extract the first element of a ShortByteString, which must be at least one Word16.
 -- An exception will be thrown in the case of an empty ShortByteString.
 head :: HasCallStack => ShortByteString -> Word16