alibaba--zvec/scripts/generate_standard_tokenizer_unicode.py

#!/usr/bin/env python3
#
# Generate Unicode lookup tables for the standard tokenizer.
#
# Usage:
#   scripts/generate_standard_tokenizer_unicode.py \
#       --ucd-dir /path/to/Public/17.0.0/ucd \
#       --out src/db/index/column/fts_column/tokenizer/standard_tokenizer_unicode.inc
from __future__ import annotations

import argparse
import pathlib

WORD_BREAK_CLASSES = {
    "ALetter": "ALetter",
    "CR": "CR",
    "Double_Quote": "DoubleQuote",
    "Extend": "Extend",
    "ExtendNumLet": "ExtendNumLet",
    "Format": "Format",
    "Hebrew_Letter": "HebrewLetter",
    "Katakana": "Katakana",
    "LF": "LF",
    "MidLetter": "MidLetter",
    "MidNum": "MidNum",
    "MidNumLet": "MidNumLet",
    "Newline": "Newline",
    "Numeric": "Numeric",
    "Regional_Indicator": "RegionalIndicator",
    "Single_Quote": "SingleQuote",
    "WSegSpace": "WSegSpace",
    "ZWJ": "ZWJ",
}

SCRIPT_CLASSES = {
    "Han": "Ideographic",
    "Hangul": "Hangul",
    "Hiragana": "Hiragana",
}


LINE_BREAK_COMPLEX_CONTEXT = {
    "Complex_Context",
    "SA",
}


def parse_codepoint_range(field):
    if ".." in field:
        start, end = field.split("..", 1)
        return int(start, 16), int(end, 16)
    cp = int(field, 16)
    return cp, cp


def parse_property_file(path, accepted_properties):
    ranges = []
    with path.open("r", encoding="utf-8") as fin:
        for raw_line in fin:
            line = raw_line.split("#", 1)[0].strip()
            if not line:
                continue
            fields = [field.strip() for field in line.split(";")]
            if len(fields) < 2:
                continue
            prop = fields[1]
            if prop not in accepted_properties:
                continue
            start, end = parse_codepoint_range(fields[0])
            ranges.append((start, end, accepted_properties[prop]))
    return merge_class_ranges(ranges)


def parse_range_properties(path, accepted_properties):
    ranges = []
    with path.open("r", encoding="utf-8") as fin:
        for raw_line in fin:
            line = raw_line.split("#", 1)[0].strip()
            if not line:
                continue
            fields = [field.strip() for field in line.split(";")]
            if len(fields) < 2:
                continue
            if fields[1] not in accepted_properties:
                continue
            ranges.append(parse_codepoint_range(fields[0]))
    return merge_ranges(ranges)


def merge_class_ranges(ranges):
    merged = []
    for start, end, cls in sorted(ranges):
        if merged and merged[-1][2] == cls and merged[-1][1] + 1 == start:
            merged[-1] = (merged[-1][0], end, cls)
        else:
            merged.append((start, end, cls))
    return merged


def merge_ranges(ranges):
    merged = []
    for start, end in sorted(ranges):
        if merged and merged[-1][1] + 1 >= start:
            merged[-1] = (merged[-1][0], max(merged[-1][1], end))
        else:
            merged.append((start, end))
    return merged


def parse_binary_property(path, property_name):
    return parse_range_properties(path, {property_name})


def write_class_table(out, name, ranges):
    out.write(f"constexpr UnicodeClassRange {name}[] = {{\n")
    for start, end, cls in ranges:
        out.write(f"    {{0x{start:04X}, 0x{end:04X}, WordBreakClass::{cls}}},\n")
    out.write("};\n\n")


def write_range_table(out, name, ranges):
    out.write(f"constexpr UnicodeRange {name}[] = {{\n")
    for start, end in ranges:
        out.write(f"    {{0x{start:04X}, 0x{end:04X}}},\n")
    out.write("};\n\n")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--ucd-dir", required=True, type=pathlib.Path)
    parser.add_argument("--out", required=True, type=pathlib.Path)
    args = parser.parse_args()

    word_break_path = args.ucd_dir / "auxiliary" / "WordBreakProperty.txt"
    emoji_data_path = args.ucd_dir / "emoji" / "emoji-data.txt"
    line_break_path = args.ucd_dir / "LineBreak.txt"
    scripts_path = args.ucd_dir / "Scripts.txt"

    word_break_ranges = parse_property_file(word_break_path, WORD_BREAK_CLASSES)
    script_ranges = parse_property_file(scripts_path, SCRIPT_CLASSES)
    extended_pictographic_ranges = parse_binary_property(
        emoji_data_path, "Extended_Pictographic"
    )
    emoji_modifier_base_ranges = parse_binary_property(
        emoji_data_path, "Emoji_Modifier_Base"
    )
    emoji_modifier_ranges = parse_binary_property(emoji_data_path, "Emoji_Modifier")
    line_break_complex_context_ranges = parse_range_properties(
        line_break_path, LINE_BREAK_COMPLEX_CONTEXT
    )

    args.out.parent.mkdir(parents=True, exist_ok=True)
    with args.out.open("w", encoding="utf-8", newline="\n") as fout:
        fout.write("// Generated by scripts/generate_standard_tokenizer_unicode.py\n")
        fout.write(
            "// Source: Unicode 17.0.0 WordBreakProperty, emoji-data, LineBreak, Scripts.\n"
        )
        fout.write(
            "// Derived from Unicode data files licensed under Unicode License V3; see NOTICE.\n"
        )
        fout.write("// Do not edit by hand.\n\n")
        fout.write("// clang-format off\n\n")
        write_class_table(fout, "kWordBreakRanges", word_break_ranges)
        write_class_table(fout, "kScriptClassRanges", script_ranges)
        write_range_table(
            fout, "kExtendedPictographicRanges", extended_pictographic_ranges
        )
        write_range_table(fout, "kEmojiModifierBaseRanges", emoji_modifier_base_ranges)
        write_range_table(fout, "kEmojiModifierRanges", emoji_modifier_ranges)
        write_range_table(
            fout, "kLineBreakComplexContextRanges", line_break_complex_context_ranges
        )
        fout.write("// clang-format on\n")


if __name__ == "__main__":
    main()