feat: DecodeResult.terminate

Author: Qiu233

Binary/Basic.lean

@@ -24,7 +24,7 @@ instance : ToString DecodeError where
 inductive DecodeResult (α) where
   | success (x : α) (k : Decoder)
   | error (err : DecodeError) (cur : Decoder)
-  | pending (fn : ByteArray → DecodeResult α)
+  | pending (fn : Option ByteArray → DecodeResult α)
 deriving Inhabited
 
 @[always_inline]
@@ -51,6 +51,7 @@ def DecodeResult.toExceptString : DecodeResult α → Except String α
 @[expose]
 def Get (α : Type) : Type := Decoder → (DecodeResult α)
 
+/-- The bytes in the current inner buffer, **not** considering pending input. -/
 @[always_inline]
 def remaining : Get Nat := fun d => DecodeResult.success (d.data.size - d.offset) d
 
@@ -60,6 +61,15 @@ def DecodeResult.feed {α} (bytes : ByteArray) : DecodeResult α → DecodeResul
   | .error err k => .error err (k.append bytes)
   | .pending fn => fn bytes
 
+/--
+Immediately terminate the pending state.
+This is especially useful for terminating something like `optional (pending x)`. -/
+@[always_inline]
+def DecodeResult.terminate {α} : DecodeResult α → DecodeResult α
+  | .success x k => .success x k
+  | .error err k => .error err k
+  | .pending fn => fn none
+
 @[always_inline]
 instance : Monad Get where
   pure x := fun d => DecodeResult.success x d
@@ -191,12 +201,15 @@ Catch any `DecodeError.eoi` and recover to a pending state rather than exit with
 * ensure that `x` terminates when enough bytes are fed,
 * or define your own pending function to cache intermediate result as much as possible.
 -/
-@[specialize, always_inline]
+@[specialize]
 partial def pending (x : Get α) : Get α := do
   try x
   catch err =>
     match err with
-    | .eoi => fun d => .pending fun bytes => pending x (d.append bytes)
+    | .eoi => fun d => .pending fun bytes? =>
+      match bytes? with
+      | none => DecodeResult.mkEOI d
+      | some bytes => pending x (d.append bytes)
     | e => throw e
 
 namespace Primitive

Binary/Get.lean

@@ -11,15 +11,15 @@ public section
 def fail (msg : String) : Get α :=
   throw (.userError msg)
 
-@[always_inline, specialize]
+@[specialize]
 def many (p : Get α) : Get (Array α) := do
   let mut data := #[]
   repeat
     let some x ← optional p | break
     data := data.push x
   return data
 
-@[always_inline, specialize]
+@[inline, specialize]
 def many1 (p : Get α) : Get (Array α) := do
   let first ← p
   let rest ← many p
@@ -35,6 +35,13 @@ def shouldBeEOI (includeUnExpected : Bool := false) : Get Unit := do
     else
       fail "expected EOI"
 
+@[always_inline]
+def notFollowedBy (p : Get α) : Get Unit := fun d =>
+  match p d with
+  | .success _ _ => DecodeResult.error (.userError "unexpected lookahead") d
+  | .error _ _ => DecodeResult.success () d
+  | .pending _ => DecodeResult.error (.userError "unexpected pending lookahead") d
+
 -- TODO: refactor following definitions for performance
 
 @[inline, specialize]
@@ -153,6 +160,26 @@ instance : Decode Int8 where
 
 end
 
+public section
+
+/--
+This function **exhaustively** reads in all bytes starting from the current offset.
+The outermost caller **must** call `DecodeResult.terminate` to break from this function. -/
+def exhaust : Get ByteArray := do
+  let r ← remaining
+  let data ← get_bytes r
+  let mut rs := #[]
+  repeat
+    shrink
+    let some x ← (optional <| pending <| getThe UInt8) | break
+    let r ← remaining
+    let xs ← get_bytes r
+    rs := rs.push ⟨#[x]⟩
+    rs := rs.push xs
+  return rs.foldl (init := data) (· ++ ·)
+
+end
+
 namespace Primitive
 
 variable {ω m} [Monad m] [STWorld ω m] [MonadLiftT (ST ω) m]

Binary/UTF8/Get.lean

@@ -11,10 +11,6 @@ namespace Binary.UTF8
 private def byteToChar (b : UInt8) : Char :=
   Char.ofNat b.toNat
 
--- @[always_inline]
--- private def charToString (c : Char) : String :=
---   String.ofList [c] -- TODO: how about `Char.toString`?
-
 @[always_inline]
 private def chars_to_string (xs : Array Char) : String :=
   String.ofList xs.toList
@@ -50,10 +46,3 @@ def manyChars (p : Get Char) : Get String :=
 @[always_inline, specialize]
 def many1Chars (p : Get Char) : Get String :=
   chars_to_string <$> many1 p
-
-@[always_inline]
-def notFollowedBy (p : Get α) : Get Unit := fun d =>
-  match p d with
-  | .success _ _ => DecodeResult.error (.userError "unexpected lookahead") d
-  | .error _ _ => DecodeResult.success () d
-  | .pending _ => DecodeResult.error (.userError "unexpected pending lookahead") d