Speed up regex validation with vectorcall - #1295

import gc
import importlib.util
import json
import os
import sys
from pathlib import Path
from typing import Annotated, ClassVar, Generic, Literal, TypeVar

import msgspec
import pyperf
from msgspec import Meta


def valid_text(value):
    if isinstance(value, str):
        if '\x00' in value:
            raise ValueError('NUL is not permitted')
        try:
            value.encode('utf-8')
        except UnicodeEncodeError:
            raise ValueError('Unicode scalar values required') from None

Source = Literal['geonames', 'ourairports', 'local']

Identifier = Annotated[str, Meta(min_length=1, max_length=160, pattern=r'^[a-z]+:[a-z]+:[A-Za-z0-9_.-]+$')]

Name = Annotated[str, Meta(min_length=1, max_length=200)]

Population = Annotated[int, Meta(ge=0, le=9223372036854775807)]

Latitude = Annotated[float, Meta(ge=-90, le=90)]

Longitude = Annotated[float, Meta(ge=-180, le=180)]

class Record(msgspec.Struct, forbid_unknown_fields=True):
    kind: ClassVar[str]
    id: Identifier
    source: Source
    name: Name

    def __post_init__(self):
        for field in self.__struct_fields__:
            valid_text(getattr(self, field))
        if not self.id.startswith(f'{self.source}:{self.kind}:'):
            raise ValueError('id namespace does not match source and resource')

class City(Record):
    kind = 'city'
    source: Literal['geonames', 'local']
    country_id: Identifier
    division_id: Identifier | None
    latitude: Latitude
    longitude: Longitude
    population: Population
    feature_code: Literal['PPL', 'PPLA', 'PPLA2', 'PPLA3', 'PPLA4', 'PPLA5', 'PPLC']

T = TypeVar('T', bound=Record)

class Page(msgspec.Struct, Generic[T], forbid_unknown_fields=True):
    items: list[T]
    next_after_id: str | None


ROWS = json.loads(
    Path(__file__).with_name("cities100.json").read_text(encoding="utf-8")
)


assert importlib.util.find_spec("pcre2") is None
probe = msgspec.Meta(pattern=r"^[A-Z]{2}$")
assert any(
    type(ref).__module__ == "re" and type(ref).__name__ == "Pattern"
    for ref in gc.get_referents(probe)
)
identifiers = [
    row[field]
    for row in ROWS
    for field in ("id", "country_id", "division_id")
    if row[field] is not None
]
assert len(ROWS) == 100 and len(identifiers) == 299
page = {"items": ROWS, "next_after_id": ROWS[-1]["id"]}
ID = Annotated[str, msgspec.Meta(pattern=r"^[a-z]+:[a-z]+:[A-Za-z0-9_.-]+$")]
CountryCode = Annotated[str, msgspec.Meta(pattern=r"^[A-Z]{2}$")]
encode = msgspec.json.Encoder().encode
cases = {
    "page100": (msgspec.json.Decoder(Page[City]), encode(page)),
    "identifier299": (msgspec.json.Decoder(list[ID]), encode(identifiers)),
    "country16": (
        msgspec.json.Decoder(list[CountryCode]),
        encode([
            "AU", "BR", "CA", "CN", "DE", "FR", "GB", "IN",
            "JP", "MX", "NZ", "RU", "SG", "US", "ZA", "AE",
        ]),
    ),
    "no_pattern299": (msgspec.json.Decoder(list[str]), encode(identifiers)),
}
operations = {}
for name, (decoder, raw) in cases.items():
    decoded = decoder.decode(raw)
    assert json.loads(encode(decoded)) == json.loads(raw)

    def operation(read=decoder.decode, raw=raw):
        return encode(read(raw))

    operations[name] = operation

# Check search semantics before timing; use escaped Unicode examples.
edge = msgspec.json.Decoder(Annotated[str, msgspec.Meta(pattern=r"^\w+$")])
for text in ["abc", "\u03b1\u03b2\u03b3", "\u00e9", "A_9", "abc\n"]:
    assert edge.decode(encode(text)) == text
for text in ["", "x\x00", "\U0001f600", "\ud800"]:
    try:
        edge.decode(encode(text))
    except (msgspec.ValidationError, msgspec.EncodeError, UnicodeError):
        pass
    else:
        raise AssertionError(text)

if __name__ == "__main__":
    if "--validate" in sys.argv:
        print({name: len(raw) for name, (_, raw) in cases.items()})
    else:
        runner = pyperf.Runner()
        runner.metadata.update(
            cell=os.environ.get("CELL", "unspecified"),
            scope="warm JSON decode/validate/encode; no PCRE2",
        )
        for name, operation in operations.items():
            runner.bench_func(name, operation)