Merge pull request #118 from konsumlamm/bsindex

Add `bsindex` primitive
augustss · Feb 3, 2025 · 6b6986f · 6b6986f
2 parents 29d7547 + 2553d4c
commit 6b6986f
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 26 deletions.
diff --git a/lib/Data/ByteString.hs b/lib/Data/ByteString.hs
@@ -130,7 +130,7 @@ module Data.ByteString(
   useAsCStringLen,
 
   -- * I\/O with 'ByteString's
-  
+
   -- ** Standard input and output
   getLine,
   getContents,
@@ -201,7 +201,9 @@ snoc :: ByteString -> Word8 -> ByteString
 snoc bs c = append bs (pack [c])
 
 head :: ByteString -> Word8
-head = P.head . unpack
+head bs
+  | null bs = bsError "head: empty"
+  | otherwise = primBSindex bs 0
 
 tail :: ByteString -> ByteString
 tail bs | sz == 0 = bsError "tail: empty"
@@ -213,7 +215,9 @@ uncons bs | null bs = Nothing
           | otherwise = Just (head bs, tail bs)
 
 last :: ByteString -> Word8
-last = P.last . unpack
+last bs
+  | null bs = bsError "last: empty"
+  | otherwise = primBSindex bs (length bs - 1)
 
 init :: ByteString -> ByteString
 init bs | sz == 0 = bsError "init: empty"
@@ -397,10 +401,15 @@ intercalate :: ByteString -> [ByteString] -> ByteString
 intercalate s = pack . P.intercalate (unpack s) . P.map unpack
 
 index :: ByteString -> Int -> Word8
-index bs n = unpack bs !! n
+index bs i
+  | i < 0          = bsError "index: negative index"
+  | i >= length bs = bsError "index: index too large"
+  | otherwise      = primBSindex bs i
 
 indexMaybe :: ByteString -> Int -> Maybe Word8
-indexMaybe bs n = unpack bs P.!? n
+indexMaybe bs i
+  | i < 0 || i >= length bs = Nothing
+  | otherwise               = Just (primBSindex bs i)
 
 (!?) :: ByteString -> Int -> Maybe Word8
 (!?) = indexMaybe

diff --git a/lib/Data/ByteString/Internal.hs b/lib/Data/ByteString/Internal.hs
@@ -30,6 +30,8 @@ primBSlength  :: ByteString -> Int
 primBSlength  = _primitive "bslength"
 primBSsubstr  :: ByteString -> Int -> Int -> ByteString
 primBSsubstr  = _primitive "bssubstr"
+primBSindex   :: ByteString -> Int -> Word8
+primBSindex   = _primitive "bsindex"
 
 -----------------------------------------
 
@@ -73,8 +75,8 @@ append = primBSappend
 
 substr :: ByteString -> Int -> Int -> ByteString
 substr bs offs len
-  | offs < 0 || offs > sz     = bsError "substr bad offset"
-  | len < 0  || len > sz-offs = bsError "substr bad length"
+  | offs < 0 || offs > sz     = bsError "substr: bad offset"
+  | len < 0  || len > sz-offs = bsError "substr: bad length"
   | otherwise = primBSsubstr bs offs len
   where sz = length bs
 

diff --git a/src/MicroHs/Translate.hs b/src/MicroHs/Translate.hs
@@ -157,6 +157,7 @@ primTable = [
   ("bsunpack", _primitive "bsunpack"),
   ("bslength", _primitive "bslength"),
   ("bssubstr", _primitive "bssubstr"),
+  ("bsindex", _primitive "bsindex"),
   ("fromUTF8", _primitive "fromUTF8"),
   ("toUTF8", _primitive "toUTF8"),
   ("headUTF8", _primitive "headUTF8"),

diff --git a/src/runtime/eval.c b/src/runtime/eval.c
@@ -266,7 +266,7 @@ enum node_tag { T_FREE, T_IND, T_AP, T_INT, T_DBL, T_PTR, T_FUNPTR, T_FORPTR, T_
                 T_IO_CCALL, T_IO_GC, T_DYNSYM,
                 T_NEWCASTRINGLEN, T_PEEKCASTRING, T_PEEKCASTRINGLEN,
                 T_BSAPPEND, T_BSAPPEND3, T_BSEQ, T_BSNE, T_BSLT, T_BSLE, T_BSGT, T_BSGE, T_BSCMP,
-                T_BSPACK, T_BSUNPACK, T_BSLENGTH, T_BSSUBSTR,
+                T_BSPACK, T_BSUNPACK, T_BSLENGTH, T_BSSUBSTR, T_BSINDEX,
                 T_BSFROMUTF8, T_BSTOUTF8, T_BSHEADUTF8,  T_BSTAILUTF8,
                 T_BSAPPENDDOT,
                 T_LAST_TAG,
@@ -768,20 +768,21 @@ struct {
   { "fread", T_FREAD},
 #endif  /* WANT_FLOAT */
 
-  { "bs++", T_BSAPPEND},
-  { "bs++.", T_BSAPPENDDOT},
-  { "bs+++", T_BSAPPEND3},
-  { "bs==", T_BSEQ, T_BSEQ},
-  { "bs/=", T_BSNE, T_BSNE},
-  { "bs<", T_BSLT},
-  { "bs<=", T_BSLE},
-  { "bs>", T_BSGT},
-  { "bs>=", T_BSGE},
-  { "bscmp", T_BSCMP},
-  { "bspack", T_BSPACK},
-  { "bsunpack", T_BSUNPACK},
-  { "bslength", T_BSLENGTH},
-  { "bssubstr", T_BSSUBSTR},
+  { "bs++", T_BSAPPEND },
+  { "bs++.", T_BSAPPENDDOT },
+  { "bs+++", T_BSAPPEND3 },
+  { "bs==", T_BSEQ, T_BSEQ },
+  { "bs/=", T_BSNE, T_BSNE },
+  { "bs<", T_BSLT },
+  { "bs<=", T_BSLE },
+  { "bs>", T_BSGT },
+  { "bs>=", T_BSGE },
+  { "bscmp", T_BSCMP },
+  { "bspack", T_BSPACK },
+  { "bsunpack", T_BSUNPACK },
+  { "bslength", T_BSLENGTH },
+  { "bssubstr", T_BSSUBSTR },
+  { "bsindex", T_BSINDEX },
 
   { "ord", T_I },
   { "chr", T_I },
@@ -2294,6 +2295,7 @@ printrec(BFILE *f, struct print_bits *pb, NODEPTR n, int prefix)
   case T_BSUNPACK: putsb("bsunpack", f); break;
   case T_BSLENGTH: putsb("bslength", f); break;
   case T_BSSUBSTR: putsb("bssubstr", f); break;
+  case T_BSINDEX: putsb("bsindex", f); break;
   case T_EQ: putsb("==", f); break;
   case T_NE: putsb("/=", f); break;
   case T_LT: putsb("<", f); break;
@@ -3479,6 +3481,14 @@ evali(NODEPTR an)
     SETBSTR(n, mkForPtrFree(rbs));
     RET;
 
+  case T_BSLENGTH:
+    CHECK(1);
+    xfp = evalbstr(ARG(TOP(0)));
+    POP(1);
+    n = TOP(-1);
+    SETINT(n, xfp->payload.size);
+    RET;
+
   case T_BSSUBSTR:
     CHECK(3);
     xfp = evalbstr(ARG(TOP(0)));
@@ -3489,12 +3499,13 @@ evali(NODEPTR an)
     SETBSTR(n, bssubstr(xfp, xi, yi));
     RET;
 
-  case T_BSLENGTH:
-    CHECK(1);
+  case T_BSINDEX:
+    CHECK(2);
     xfp = evalbstr(ARG(TOP(0)));
-    POP(1);
+    xi = evalint(ARG(TOP(1)));
+    POP(2);
     n = TOP(-1);
-    SETINT(n, xfp->payload.size);
+    SETINT(n, ((uint8_t *)xfp->payload.string)[xi]);
     RET;
 
   case T_RAISE: