Skip to content

Commit a80bd07

Browse files
committed
Speed up readDriveLetter 10x by introducing uncons2
1 parent 6d371f1 commit a80bd07

File tree

4 files changed

+51
-9
lines changed

4 files changed

+51
-9
lines changed

System/FilePath/Internal.hs

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -506,12 +506,14 @@ readDriveUNC bs = case unpack bs of
506506

507507
{- c:\ -}
508508
readDriveLetter :: STRING -> Maybe (FILEPATH, FILEPATH)
509-
readDriveLetter bs = case unpack bs of
510-
(x:c:y:xs)
511-
| c == _colon && isLetter x && isPathSeparator y -> Just $ addSlash (pack [x,_colon]) (pack (y:xs))
512-
(x:c:xs)
513-
| c == _colon && isLetter x -> Just (pack [x,_colon], pack xs)
514-
_ -> Nothing
509+
readDriveLetter bs = case uncons2 bs of
510+
Nothing -> Nothing
511+
Just (x, c, ys)
512+
| isLetter x, c == _colon -> Just $ case uncons ys of
513+
Just (y, _)
514+
| isPathSeparator y -> addSlash (pack [x,_colon]) ys
515+
_ -> (pack [x,_colon], ys)
516+
| otherwise -> Nothing
515517

516518
{- \\sharename\ -}
517519
readDriveShare :: STRING -> Maybe (FILEPATH, FILEPATH)
@@ -1147,6 +1149,10 @@ stripSuffix xs ys = reverse P.<$> stripPrefix (reverse xs) (reverse ys)
11471149
unsnoc :: [a] -> Maybe ([a], a)
11481150
unsnoc = L.foldr (\x -> Just . maybe ([], x) (first (x :))) Nothing
11491151

1152+
uncons2 :: [a] -> Maybe (a, a, [a])
1153+
uncons2 [] = Nothing
1154+
uncons2 [_] = Nothing
1155+
uncons2 (x : y : zs) = Just (x, y, zs)
11501156

11511157
_period, _quotedbl, _backslash, _slash, _question, _U, _N, _C, _colon, _semicolon, _US, _less, _greater, _bar, _asterisk, _nul, _space, _underscore :: Char
11521158
_period = '.'

System/OsPath/Data/ByteString/Short.hs

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
{-# LANGUAGE MultiWayIf #-}
12
{-# LANGUAGE NoImplicitPrelude #-}
23
-- |
34
-- Module : System.OsPath.Data.ByteString.Short
@@ -81,6 +82,7 @@ module System.OsPath.Data.ByteString.Short (
8182
last,
8283
tail,
8384
uncons,
85+
uncons2,
8486
head,
8587
init,
8688
unsnoc,
@@ -173,3 +175,17 @@ module System.OsPath.Data.ByteString.Short (
173175
) where
174176

175177
import Data.ByteString.Short.Internal
178+
import System.OsPath.Data.ByteString.Short.Internal
179+
180+
import Prelude (Maybe(..), Ord(..), Num(..), ($), otherwise)
181+
import Data.Word (Word8)
182+
183+
uncons2 :: ShortByteString -> Maybe (Word8, Word8, ShortByteString)
184+
uncons2 = \sbs ->
185+
let l = length sbs
186+
nl = l - 2
187+
in if | l <= 1 -> Nothing
188+
| otherwise -> let h = indexWord8Array (asBA sbs) 0
189+
h' = indexWord8Array (asBA sbs) 1
190+
t = create nl $ \mba -> copyByteArray (asBA sbs) 1 mba 0 nl
191+
in Just (h, h', t)

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ module System.OsPath.Data.ByteString.Short.Internal where
2020

2121
import Control.Monad.ST
2222
import Control.Exception (assert, throwIO)
23+
import Data.Bits (Bits(..))
2324
import Data.ByteString.Short.Internal (ShortByteString(..), length)
2425
#if !MIN_VERSION_base(4,11,0)
2526
import Data.Semigroup
@@ -284,15 +285,21 @@ writeWord16Array (MBA# mba#) (I# i#) (W16# w#) =
284285
ST (\s -> case writeWord8Array# mba# (i# +# 1#) msb# s of
285286
s' -> (# s', () #))
286287

288+
indexWord8Array :: BA
289+
-> Int -- ^ Word8 index
290+
-> Word8
291+
indexWord8Array (BA# ba#) (I# i#) = W8# (indexWord8Array# ba# i#)
292+
287293
-- | This isn't strictly Word16 array read. Instead it's two Word8 array reads
288294
-- to avoid endianness issues due to primops doing automatic alignment based
289295
-- on host platform. We expect the byte array to be LE always.
290296
indexWord16Array :: BA
291297
-> Int -- ^ Word8 index (not Word16)
292298
-> Word16
293-
indexWord16Array (BA# ba#) (I# i#) =
294-
case (# indexWord8Array# ba# i#, indexWord8Array# ba# (i# +# 1#) #) of
295-
(# lsb#, msb# #) -> W16# (decodeWord16LE# (# lsb#, msb# #))
299+
indexWord16Array ba i = fromIntegral lsb .|. (fromIntegral msb `shiftL` 8)
300+
where
301+
lsb = indexWord8Array ba i
302+
msb = indexWord8Array ba (i + 1)
296303

297304
#if !MIN_VERSION_base(4,16,0)
298305

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ module System.OsPath.Data.ByteString.Short.Word16 (
4646
last,
4747
tail,
4848
uncons,
49+
uncons2,
4950
head,
5051
init,
5152
unsnoc,
@@ -260,6 +261,18 @@ uncons = \(assertEven -> sbs) ->
260261
t = create nl $ \mba -> copyByteArray (asBA sbs) 2 mba 0 nl
261262
in Just (h, t)
262263

264+
-- | /O(n)/ Extract first two elements and the rest of a ByteString,
265+
-- returning Nothing if it is shorter than two elements.
266+
uncons2 :: ShortByteString -> Maybe (Word16, Word16, ShortByteString)
267+
uncons2 = \(assertEven -> sbs) ->
268+
let l = BS.length sbs
269+
nl = l - 4
270+
in if | l <= 2 -> Nothing
271+
| otherwise -> let h = indexWord16Array (asBA sbs) 0
272+
h' = indexWord16Array (asBA sbs) 2
273+
t = create nl $ \mba -> copyByteArray (asBA sbs) 4 mba 0 nl
274+
in Just (h, h', t)
275+
263276
-- | /O(1)/ Extract the first element of a ShortByteString, which must be at least one Word16.
264277
-- An exception will be thrown in the case of an empty ShortByteString.
265278
head :: HasCallStack => ShortByteString -> Word16

0 commit comments

Comments
 (0)