Enum case matching by remaining buffer length (#26)

This PR adds support of matching enum case using remaining buffer
length. This matching method is exclusive with byte matching.

The changes (implementation, test, and specs) are reviewed by the owner.

Author: FlickerSoul

Sources/BinaryParseKit/BinaryParseKit.swift

@@ -401,6 +401,44 @@ public macro match(bytes: [UInt8]) = #externalMacro(
     type: "EmptyPeerMacro",
 )
 
+/// Matches when the remaining buffer bytes equals the specified length exactly.
+///
+/// Use this macro to match cases based on remaining data size rather than specific byte patterns.
+/// The buffer is NOT consumed, allowing subsequent parsing operations to read the data.
+///
+/// - Parameter length: The exact number of remaining bytes to match
+///
+/// - Note: An enum using `@ParseEnum` can only use byte-based matching OR length-based matching,
+///   but not both. `@matchDefault` is allowed with either strategy. That is, it's not allowed to
+///   use a mix of `@match(byte:)`/`@match(bytes:)`/`@matchAndTake(byte:)` and `@match(length:)`
+///   in the same enum.
+///
+///
+/// Example:
+/// ```swift
+/// @ParseEnum
+/// enum VariableSizeData {
+///     @match(length: 4)
+///     @parse(endianness: .big)
+///     case shortPayload(UInt32)
+///
+///     @match(length: 8)
+///     @parse(endianness: .big)
+///     case longPayload(UInt64)
+///
+///     @matchDefault
+///     case unknown
+/// }
+///
+/// let data = try VariableSizeData(parsing: Data([0x12, 0x34, 0x56, 0x78]))
+/// // data == .shortPayload(0x12345678) because buffer has exactly 4 bytes
+/// ```
+@attached(peer)
+public macro match(length: Int) = #externalMacro(
+    module: "BinaryParseKitMacros",
+    type: "EmptyPeerMacro",
+)
+
 /// Matches and consumes bytes from the buffer using the enum's raw value.
 ///
 /// Use this macro for ``Matchable`` enums where each case's raw value

Sources/BinaryParseKit/Utils/ParsingUtils.swift

@@ -7,7 +7,7 @@
 import BinaryParsing
 
 /// Matches the given bytes in the input parser span.
-/// - Warning: This function is used to `@parseEnum` macro and should not be used directly.
+/// - Warning: This function is used by `@ParseEnum` macro and should not be used directly.
 @inlinable
 public func __match(_ bytes: borrowing [UInt8], in input: inout BinaryParsing.ParserSpan) -> Bool {
     if bytes.isEmpty { return true }
@@ -22,6 +22,13 @@ public func __match(_ bytes: borrowing [UInt8], in input: inout BinaryParsing.Pa
     return toMatch == bytes
 }
 
+/// Matches when the remaining bytes in the input parser span equals the specified length.
+/// - Warning: This function is used by `@ParseEnum` macro and should not be used directly.
+@inlinable
+public func __match(length: Int, in input: borrowing BinaryParsing.ParserSpan) -> Bool {
+    input.count == length
+}
+
 /// Asserts that the given type conforms to `Parsable`.
 /// - Warning: This function is used to `@parse` macro and should not be used directly.
 @inlinable

…s/ParseEnum/ConsructParseEnumMacro.swift …/ParseEnum/ConstructParseEnumMacro.swiftSources/BinaryParseKitMacros/Macros/ParseEnum/ConsructParseEnumMacro.swift renamed to Sources/BinaryParseKitMacros/Macros/ParseEnum/ConstructParseEnumMacro.swift Sources/BinaryParseKitMacros/Macros/ParseEnum/ConsructParseEnumMacro.swift renamed to Sources/BinaryParseKitMacros/Macros/ParseEnum/ConstructParseEnumMacro.swift

@@ -1,5 +1,5 @@
 //
-//  ConsructParseEnumMacro.swift
+//  ConstructParseEnumMacro.swift
 //  BinaryParseKit
 //
 //  Created by Larry Zeng on 7/26/25.
@@ -43,13 +43,36 @@ public struct ConstructEnumParseMacro: ExtensionMacro {
                     "\(accessorInfo.parsingAccessor) init(parsing span: inout \(raw: Constants.BinaryParsing.parserSpan)) throws(\(raw: Constants.BinaryParsing.thrownParsingError))",
                 ) {
                     for caseParseInfo in parseInfo.caseParseInfo {
-                        let toBeMatched = caseParseInfo.bytesToMatch(of: type)
+                        // Generate the match condition based on match type
+                        let matchCondition = try ConditionElementListSyntax {
+                            if let matchLength = caseParseInfo.lengthToMatch() {
+                                // Length-based matching: use __match(length:in:) with borrowing span
+                                ExprSyntax(
+                                    "\(raw: Constants.UtilityFunctions.matchLength)(length: \(matchLength), in: span)",
+                                )
+                            } else if let toBeMatched = caseParseInfo.bytesToMatch(of: type) {
+                                // Byte-array-based matching: use __matchBytes with inout span
+                                ExprSyntax("\(raw: Constants.UtilityFunctions.matchBytes)(\(toBeMatched), in: &span)")
+                            } else if caseParseInfo.matchAction.target.isDefaultMatch {
+                                // Default matching: always true
+                                ExprSyntax("true")
+                            } else {
+                                // Otherwise, it's a failure on our side
+                                throw ParseEnumMacroError.unexpectedError(
+                                    description: "Failed to obtain matching bytes for \(caseParseInfo.caseElementName)",
+                                )
+                            }
+                        }
 
-                        try IfExprSyntax(
-                            "if \(raw: Constants.UtilityFunctions.matchBytes)(\(toBeMatched), in: &span)",
-                        ) {
+                        try IfExprSyntax("if \(matchCondition)") {
                             if caseParseInfo.matchAction.matchPolicy == .matchAndTake {
-                                "try span.seek(toRelativeOffset: \(toBeMatched).count)"
+                                if let toBeMatched = caseParseInfo.bytesToMatch(of: type) {
+                                    "try span.seek(toRelativeOffset: \(toBeMatched).count)"
+                                } else {
+                                    throw ParseEnumMacroError.unexpectedError(
+                                        description: "Failed to obtain matching bytes for \(caseParseInfo.caseElementName) when taking",
+                                    )
+                                }
                             }
 
                             var arguments: OrderedDictionary<TokenSyntax, EnumCaseParameterParseInfo> = [:]
@@ -179,16 +202,29 @@ public struct ConstructEnumParseMacro: ExtensionMacro {
                                     }
                                 })
 
-                            let caseCodeBlock = CodeBlockItemListSyntax {
-                                let toBeMatched = caseParseInfo.bytesToMatch(of: type)
+                            let caseCodeBlock = try CodeBlockItemListSyntax {
+                                let bytesTakenInMatching = context.makeUniqueName("bytesTakenInMatching")
+
+                                if caseParseInfo.matchAction.target.isLengthMatch
+                                    || caseParseInfo.matchAction.target.isDefaultMatch {
+                                    // For length-based matching, use empty bytes array (similar to matchDefault)
+                                    "let \(bytesTakenInMatching): [UInt8] = []"
+                                } else if let bytesToMatch = caseParseInfo.bytesToMatch(of: type) {
+                                    "let \(bytesTakenInMatching): [UInt8] = \(bytesToMatch)"
+                                } else {
+                                    throw ParseEnumMacroError.unexpectedError(
+                                        description: "Failed to obtain matching bytes for \(caseParseInfo.caseElementName)",
+                                    )
+                                }
+
                                 let matchPolicy = caseParseInfo.matchAction.matchPolicy
 
                                 let fields = ArrayExprSyntax(elements: generatePrintableFields(parseSkipMacroInfo))
 
                                 #"""
                                 return .enum(
                                     .init(
-                                        bytes: \#(toBeMatched),
+                                        bytes: \#(bytesTakenInMatching),
                                         parseType: .\#(raw: matchPolicy),
                                         fields: \#(fields),
                                     )

Sources/BinaryParseKitMacros/Macros/ParseEnum/EnumCaseParseInfo.swift

@@ -45,26 +45,92 @@ enum EnumParseAction {
     case skip(SkipMacroInfo)
 }
 
+enum EnumMatchTarget {
+    // @match() -> .bytes(nil)
+    // @match(byte: byte) -> .byte([byte])
+    // @match(bytes: bytes) -> .bytes(bytes)
+    // @matchAndTake() -> .bytes(nil)
+    // @matchAndTake(byte: byte) -> .byte([byte])
+    // @matchAndTake(bytes: bytes) -> .bytes(bytes)
+    case bytes(ExprSyntax?)
+    // @match(length: length) -> .length(length)
+    case length(ExprSyntax)
+    // @matchDefault
+    case `default`
+
+    var matchBytes: ExprSyntax?? {
+        if case let .bytes(bytes) = self {
+            bytes
+        } else {
+            nil
+        }
+    }
+
+    var matchLength: ExprSyntax? {
+        if case let .length(length) = self {
+            length
+        } else {
+            nil
+        }
+    }
+
+    var isLengthMatch: Bool {
+        if case .length = self {
+            true
+        } else {
+            false
+        }
+    }
+
+    var isByteMatch: Bool {
+        if case .bytes = self {
+            true
+        } else {
+            false
+        }
+    }
+
+    var isDefaultMatch: Bool {
+        if case .default = self {
+            true
+        } else {
+            false
+        }
+    }
+}
+
 struct EnumCaseMatchAction {
-    let matchBytes: ExprSyntax?
+    let target: EnumMatchTarget
     let matchPolicy: EnumCaseMatchPolicy
 
     static func match(bytes: ExprSyntax?) -> EnumCaseMatchAction {
-        EnumCaseMatchAction(matchBytes: bytes, matchPolicy: .match)
+        EnumCaseMatchAction(target: .bytes(bytes), matchPolicy: .match)
+    }
+
+    static func matchLength(_ length: ExprSyntax) -> EnumCaseMatchAction {
+        EnumCaseMatchAction(target: .length(length), matchPolicy: .match)
     }
 
     static func matchDefault() -> EnumCaseMatchAction {
-        EnumCaseMatchAction(matchBytes: "[]", matchPolicy: .matchDefault)
+        EnumCaseMatchAction(target: .default, matchPolicy: .matchDefault)
     }
 
     static func matchAndTake(bytes: ExprSyntax?) -> EnumCaseMatchAction {
-        EnumCaseMatchAction(matchBytes: bytes, matchPolicy: .matchAndTake)
+        EnumCaseMatchAction(target: .bytes(bytes), matchPolicy: .matchAndTake)
     }
 
     static func parseMatch(from attribute: AttributeSyntax) throws(ParseEnumMacroError) -> EnumCaseMatchAction {
         let arguments = attribute.arguments?.as(LabeledExprListSyntax.self)
-        let bytes = try parseBytesArgument(in: arguments, at: 0)
 
+        // Check if this is @match(length:)
+        if let args = arguments,
+           let firstArg = args.first,
+           firstArg.label?.text == "length" {
+            return .matchLength(firstArg.expression)
+        }
+
+        // Otherwise, parse as byte-based match
+        let bytes = try parseBytesArgument(in: arguments, at: 0)
         return .match(bytes: bytes)
     }
 
@@ -117,22 +183,35 @@ struct EnumCaseParseInfo {
     let matchAction: EnumCaseMatchAction
     let parseActions: [EnumParseAction]
     let caseElementName: TokenSyntax
-
-    init(matchAction: EnumCaseMatchAction, parseActions: [EnumParseAction], caseElementName: TokenSyntax) {
+    let source: Syntax
+
+    init(
+        matchAction: EnumCaseMatchAction,
+        parseActions: [EnumParseAction],
+        caseElementName: TokenSyntax,
+        source: Syntax,
+    ) {
         self.matchAction = matchAction
         self.parseActions = parseActions
         self.caseElementName = caseElementName.trimmed
+        self.source = source
     }
 
-    func bytesToMatch(of type: some TypeSyntaxProtocol) -> ExprSyntax {
-        if let matchBytes = matchAction.matchBytes {
-            matchBytes
-        } else {
-            ExprSyntax(
-                "(\(type).\(caseElementName) as any \(raw: Constants.Protocols.matchableProtocol)).bytesToMatch()",
-            )
+    func bytesToMatch(of type: some TypeSyntaxProtocol) -> ExprSyntax? {
+        matchAction.target.matchBytes.map { matchBytes in
+            if let matchBytes {
+                matchBytes
+            } else {
+                ExprSyntax(
+                    "(\(type).\(caseElementName) as any \(raw: Constants.Protocols.matchableProtocol)).bytesToMatch()",
+                )
+            }
         }
     }
+
+    func lengthToMatch() -> ExprSyntax? {
+        matchAction.target.matchLength
+    }
 }
 
 struct EnumParseInfo {

Sources/BinaryParseKitMacros/Macros/ParseEnum/ParseEnumCase.swift

@@ -18,6 +18,8 @@ class ParseEnumCase: SyntaxVisitor {
     private(set) var parsedInfo: EnumParseInfo?
 
     private var hasMatchDefault = false
+    private var hasByteMatch = false
+    private var hasLengthMatch = false
 
     private var errors: [Diagnostic] = []
 
@@ -107,6 +109,29 @@ class ParseEnumCase: SyntaxVisitor {
                     ),
                 )
             }
+
+            // Track matching strategy for mutual exclusivity validation
+            if matchAction.target.isLengthMatch {
+                if hasByteMatch {
+                    errors.append(
+                        .init(
+                            node: node,
+                            message: ParseEnumMacroError.mixedMatchingStrategies,
+                        ),
+                    )
+                }
+                hasLengthMatch = true
+            } else if matchAction.target.isByteMatch {
+                if hasLengthMatch {
+                    errors.append(
+                        .init(
+                            node: node,
+                            message: ParseEnumMacroError.mixedMatchingStrategies,
+                        ),
+                    )
+                }
+                hasByteMatch = true
+            }
         }
 
         for currentCaseElement in currentCaseElements {
@@ -120,6 +145,7 @@ class ParseEnumCase: SyntaxVisitor {
                     matchAction: matchAction,
                     parseActions: enumParseActions,
                     caseElementName: currentCaseElement.name,
+                    source: Syntax(node),
                 ),
             )
         }

Sources/BinaryParseKitMacros/Macros/ParseEnum/ParseEnumMacroError.swift

@@ -15,6 +15,7 @@ enum ParseEnumMacroError: Error, DiagnosticMessage {
     case defaultCaseMustBeLast
     case onlyOneMatchDefaultAllowed
     case matchDefaultShouldBeLast
+    case mixedMatchingStrategies
     case unexpectedError(description: String)
 
     var message: String {
@@ -29,6 +30,8 @@ enum ParseEnumMacroError: Error, DiagnosticMessage {
         case .defaultCaseMustBeLast: "The `matchDefault` case must be the last case in the enum."
         case .onlyOneMatchDefaultAllowed: "Only one `matchDefault` case is allowed in a enum."
         case .matchDefaultShouldBeLast: "The `matchDefault` case should be the last case in the enum."
+        case .mixedMatchingStrategies:
+            "An enum cannot mix byte-based matching (@match, @match(byte:), @match(bytes:), @matchAndTake) with length-based matching (@match(length:))."
         case let .unexpectedError(description: description):
             "Unexpected error: \(description)"
         }

Sources/BinaryParseKitMacros/Macros/Supports/Constants.swift

@@ -41,6 +41,7 @@ extension Constants {
 extension Constants {
     enum UtilityFunctions {
         static let matchBytes = PackageMember(name: "__match")
+        static let matchLength = PackageMember(name: "__match")
         static let assertParsable = PackageMember(name: "__assertParsable")
         static let assertSizedParsable = PackageMember(name: "__assertSizedParsable")
         static let assertEndianParsable = PackageMember(name: "__assertEndianParsable")