Use 1 byte per type/symbol tag (#19735)

This is a small incremental improvement for fixed format cache. I am
adding a dedicated write/read functions for tags (i.e. integers in 0-255
range). I propose to exclusively use these functions for type tags
(hence the name), and still use regular `write_int()`/`read_int()` for
integers that are "accidentally small" (like argument kinds etc). In a
separate PR I will change regular `int` format to be more progressive
(e.g. only use 1 byte if an integer happens to be small). I also change
the terminology from "marker" to "tag", as this is a more common name
for this concept.

Note we can probably use `mypy_extensions.u8` for type tags. If there is
a desire for this, I can switch to it (either in this or a separate PR).

Author: ilevkivskyi

mypy/cache.py

@@ -10,10 +10,12 @@
         read_float as read_float,
         read_int as read_int,
         read_str as read_str,
+        read_tag as read_tag,
         write_bool as write_bool,
         write_float as write_float,
         write_int as write_int,
         write_str as write_str,
+        write_tag as write_tag,
     )
 except ImportError:
     # TODO: temporary, remove this after we publish mypy-native on PyPI.
@@ -32,6 +34,12 @@ def read_int(data: Buffer) -> int:
         def write_int(data: Buffer, value: int) -> None:
             raise NotImplementedError
 
+        def read_tag(data: Buffer) -> int:
+            raise NotImplementedError
+
+        def write_tag(data: Buffer, value: int) -> None:
+            raise NotImplementedError
+
         def read_str(data: Buffer) -> str:
             raise NotImplementedError
 
@@ -59,37 +67,37 @@ def write_float(data: Buffer, value: float) -> None:
 LITERAL_NONE: Final = 6
 
 
-def read_literal(data: Buffer, marker: int) -> int | str | bool | float:
-    if marker == LITERAL_INT:
+def read_literal(data: Buffer, tag: int) -> int | str | bool | float:
+    if tag == LITERAL_INT:
         return read_int(data)
-    elif marker == LITERAL_STR:
+    elif tag == LITERAL_STR:
         return read_str(data)
-    elif marker == LITERAL_BOOL:
+    elif tag == LITERAL_BOOL:
         return read_bool(data)
-    elif marker == LITERAL_FLOAT:
+    elif tag == LITERAL_FLOAT:
         return read_float(data)
-    assert False, f"Unknown literal marker {marker}"
+    assert False, f"Unknown literal tag {tag}"
 
 
 def write_literal(data: Buffer, value: int | str | bool | float | complex | None) -> None:
     if isinstance(value, bool):
-        write_int(data, LITERAL_BOOL)
+        write_tag(data, LITERAL_BOOL)
         write_bool(data, value)
     elif isinstance(value, int):
-        write_int(data, LITERAL_INT)
+        write_tag(data, LITERAL_INT)
         write_int(data, value)
     elif isinstance(value, str):
-        write_int(data, LITERAL_STR)
+        write_tag(data, LITERAL_STR)
         write_str(data, value)
     elif isinstance(value, float):
-        write_int(data, LITERAL_FLOAT)
+        write_tag(data, LITERAL_FLOAT)
         write_float(data, value)
     elif isinstance(value, complex):
-        write_int(data, LITERAL_COMPLEX)
+        write_tag(data, LITERAL_COMPLEX)
         write_float(data, value.real)
         write_float(data, value.imag)
     else:
-        write_int(data, LITERAL_NONE)
+        write_tag(data, LITERAL_NONE)
 
 
 def read_int_opt(data: Buffer) -> int | None:

mypy/nodes.py

@@ -28,6 +28,7 @@
     read_str_list,
     read_str_opt,
     read_str_opt_list,
+    read_tag,
     write_bool,
     write_int,
     write_int_list,
@@ -37,6 +38,7 @@
     write_str_list,
     write_str_opt,
     write_str_opt_list,
+    write_tag,
 )
 from mypy.options import Options
 from mypy.util import is_sunder, is_typeshed_file, short_type
@@ -417,7 +419,7 @@ def deserialize(cls, data: JsonDict) -> MypyFile:
         return tree
 
     def write(self, data: Buffer) -> None:
-        write_int(data, MYPY_FILE)
+        write_tag(data, MYPY_FILE)
         write_str(data, self._fullname)
         self.names.write(data, self._fullname)
         write_bool(data, self.is_stub)
@@ -427,7 +429,7 @@ def write(self, data: Buffer) -> None:
 
     @classmethod
     def read(cls, data: Buffer) -> MypyFile:
-        assert read_int(data) == MYPY_FILE
+        assert read_tag(data) == MYPY_FILE
         tree = MypyFile([], [])
         tree._fullname = read_str(data)
         tree.names = SymbolTable.read(data)
@@ -711,7 +713,7 @@ def deserialize(cls, data: JsonDict) -> OverloadedFuncDef:
         return res
 
     def write(self, data: Buffer) -> None:
-        write_int(data, OVERLOADED_FUNC_DEF)
+        write_tag(data, OVERLOADED_FUNC_DEF)
         write_int(data, len(self.items))
         for item in self.items:
             item.write(data)
@@ -1022,7 +1024,7 @@ def deserialize(cls, data: JsonDict) -> FuncDef:
         return ret
 
     def write(self, data: Buffer) -> None:
-        write_int(data, FUNC_DEF)
+        write_tag(data, FUNC_DEF)
         write_str(data, self._name)
         mypy.types.write_type_opt(data, self.type)
         write_str(data, self._fullname)
@@ -1134,16 +1136,16 @@ def deserialize(cls, data: JsonDict) -> Decorator:
         return dec
 
     def write(self, data: Buffer) -> None:
-        write_int(data, DECORATOR)
+        write_tag(data, DECORATOR)
         self.func.write(data)
         self.var.write(data)
         write_bool(data, self.is_overload)
 
     @classmethod
     def read(cls, data: Buffer) -> Decorator:
-        assert read_int(data) == FUNC_DEF
+        assert read_tag(data) == FUNC_DEF
         func = FuncDef.read(data)
-        assert read_int(data) == VAR
+        assert read_tag(data) == VAR
         var = Var.read(data)
         dec = Decorator(func, [], var)
         dec.is_overload = read_bool(data)
@@ -1326,7 +1328,7 @@ def deserialize(cls, data: JsonDict) -> Var:
         return v
 
     def write(self, data: Buffer) -> None:
-        write_int(data, VAR)
+        write_tag(data, VAR)
         write_str(data, self._name)
         mypy.types.write_type_opt(data, self.type)
         mypy.types.write_type_opt(data, self.setter_type)
@@ -1341,13 +1343,13 @@ def read(cls, data: Buffer) -> Var:
         v = Var(name, typ)
         setter_type: mypy.types.CallableType | None = None
         if read_bool(data):
-            assert read_int(data) == mypy.types.CALLABLE_TYPE
+            assert read_tag(data) == mypy.types.CALLABLE_TYPE
             setter_type = mypy.types.CallableType.read(data)
         v.setter_type = setter_type
         v.is_ready = False  # Override True default set in __init__
         v._fullname = read_str(data)
         read_flags(data, v, VAR_FLAGS)
-        marker = read_int(data)
+        marker = read_tag(data)
         if marker == LITERAL_COMPLEX:
             v.final_value = complex(read_float(data), read_float(data))
         elif marker != LITERAL_NONE:
@@ -1465,7 +1467,7 @@ def deserialize(cls, data: JsonDict) -> ClassDef:
         return res
 
     def write(self, data: Buffer) -> None:
-        write_int(data, CLASS_DEF)
+        write_tag(data, CLASS_DEF)
         write_str(data, self.name)
         mypy.types.write_type_list(data, self.type_vars)
         write_str(data, self.fullname)
@@ -2898,7 +2900,7 @@ def deserialize(cls, data: JsonDict) -> TypeVarExpr:
         )
 
     def write(self, data: Buffer) -> None:
-        write_int(data, TYPE_VAR_EXPR)
+        write_tag(data, TYPE_VAR_EXPR)
         write_str(data, self._name)
         write_str(data, self._fullname)
         mypy.types.write_type_list(data, self.values)
@@ -2948,7 +2950,7 @@ def deserialize(cls, data: JsonDict) -> ParamSpecExpr:
         )
 
     def write(self, data: Buffer) -> None:
-        write_int(data, PARAM_SPEC_EXPR)
+        write_tag(data, PARAM_SPEC_EXPR)
         write_str(data, self._name)
         write_str(data, self._fullname)
         self.upper_bound.write(data)
@@ -3016,7 +3018,7 @@ def deserialize(cls, data: JsonDict) -> TypeVarTupleExpr:
         )
 
     def write(self, data: Buffer) -> None:
-        write_int(data, TYPE_VAR_TUPLE_EXPR)
+        write_tag(data, TYPE_VAR_TUPLE_EXPR)
         self.tuple_fallback.write(data)
         write_str(data, self._name)
         write_str(data, self._fullname)
@@ -3026,7 +3028,7 @@ def write(self, data: Buffer) -> None:
 
     @classmethod
     def read(cls, data: Buffer) -> TypeVarTupleExpr:
-        assert read_int(data) == mypy.types.INSTANCE
+        assert read_tag(data) == mypy.types.INSTANCE
         fallback = mypy.types.Instance.read(data)
         return TypeVarTupleExpr(
             read_str(data),
@@ -3908,7 +3910,7 @@ def deserialize(cls, data: JsonDict) -> TypeInfo:
         return ti
 
     def write(self, data: Buffer) -> None:
-        write_int(data, TYPE_INFO)
+        write_tag(data, TYPE_INFO)
         self.names.write(data, self.fullname)
         self.defn.write(data)
         write_str(data, self.module_name)
@@ -3944,7 +3946,7 @@ def write(self, data: Buffer) -> None:
     @classmethod
     def read(cls, data: Buffer) -> TypeInfo:
         names = SymbolTable.read(data)
-        assert read_int(data) == CLASS_DEF
+        assert read_tag(data) == CLASS_DEF
         defn = ClassDef.read(data)
         module_name = read_str(data)
         ti = TypeInfo(names, defn, module_name)
@@ -3954,10 +3956,9 @@ def read(cls, data: Buffer) -> TypeInfo:
         ti.abstract_attributes = list(zip(attrs, statuses))
         ti.type_vars = read_str_list(data)
         ti.has_param_spec_type = read_bool(data)
-        num_bases = read_int(data)
         ti.bases = []
-        for _ in range(num_bases):
-            assert read_int(data) == mypy.types.INSTANCE
+        for _ in range(read_int(data)):
+            assert read_tag(data) == mypy.types.INSTANCE
             ti.bases.append(mypy.types.Instance.read(data))
         # NOTE: ti.mro will be set in the fixup phase based on these
         # names.  The reason we need to store the mro instead of just
@@ -3972,19 +3973,19 @@ def read(cls, data: Buffer) -> TypeInfo:
         ti._mro_refs = read_str_list(data)
         ti._promote = cast(list[mypy.types.ProperType], mypy.types.read_type_list(data))
         if read_bool(data):
-            assert read_int(data) == mypy.types.INSTANCE
+            assert read_tag(data) == mypy.types.INSTANCE
             ti.alt_promote = mypy.types.Instance.read(data)
         if read_bool(data):
-            assert read_int(data) == mypy.types.INSTANCE
+            assert read_tag(data) == mypy.types.INSTANCE
             ti.declared_metaclass = mypy.types.Instance.read(data)
         if read_bool(data):
-            assert read_int(data) == mypy.types.INSTANCE
+            assert read_tag(data) == mypy.types.INSTANCE
             ti.metaclass_type = mypy.types.Instance.read(data)
         if read_bool(data):
-            assert read_int(data) == mypy.types.TUPLE_TYPE
+            assert read_tag(data) == mypy.types.TUPLE_TYPE
             ti.tuple_type = mypy.types.TupleType.read(data)
         if read_bool(data):
-            assert read_int(data) == mypy.types.TYPED_DICT_TYPE
+            assert read_tag(data) == mypy.types.TYPED_DICT_TYPE
             ti.typeddict_type = mypy.types.TypedDictType.read(data)
         read_flags(data, ti, TypeInfo.FLAGS)
         metadata = read_str(data)
@@ -3994,7 +3995,7 @@ def read(cls, data: Buffer) -> TypeInfo:
             ti.slots = set(read_str_list(data))
         ti.deletable_attributes = read_str_list(data)
         if read_bool(data):
-            assert read_int(data) == mypy.types.TYPE_VAR_TYPE
+            assert read_tag(data) == mypy.types.TYPE_VAR_TYPE
             ti.self_type = mypy.types.TypeVarType.read(data)
         if read_bool(data):
             ti.dataclass_transform_spec = DataclassTransformSpec.read(data)
@@ -4270,7 +4271,7 @@ def deserialize(cls, data: JsonDict) -> TypeAlias:
         )
 
     def write(self, data: Buffer) -> None:
-        write_int(data, TYPE_ALIAS)
+        write_tag(data, TYPE_ALIAS)
         write_str(data, self._fullname)
         self.target.write(data)
         mypy.types.write_type_list(data, self.alias_tvars)
@@ -4890,33 +4891,33 @@ def local_definitions(
 
 
 def read_symbol(data: Buffer) -> mypy.nodes.SymbolNode:
-    marker = read_int(data)
+    tag = read_tag(data)
     # The branches here are ordered manually by type "popularity".
-    if marker == VAR:
+    if tag == VAR:
         return mypy.nodes.Var.read(data)
-    if marker == FUNC_DEF:
+    if tag == FUNC_DEF:
         return mypy.nodes.FuncDef.read(data)
-    if marker == DECORATOR:
+    if tag == DECORATOR:
         return mypy.nodes.Decorator.read(data)
-    if marker == TYPE_INFO:
+    if tag == TYPE_INFO:
         return mypy.nodes.TypeInfo.read(data)
-    if marker == OVERLOADED_FUNC_DEF:
+    if tag == OVERLOADED_FUNC_DEF:
         return mypy.nodes.OverloadedFuncDef.read(data)
-    if marker == TYPE_VAR_EXPR:
+    if tag == TYPE_VAR_EXPR:
         return mypy.nodes.TypeVarExpr.read(data)
-    if marker == TYPE_ALIAS:
+    if tag == TYPE_ALIAS:
         return mypy.nodes.TypeAlias.read(data)
-    if marker == PARAM_SPEC_EXPR:
+    if tag == PARAM_SPEC_EXPR:
         return mypy.nodes.ParamSpecExpr.read(data)
-    if marker == TYPE_VAR_TUPLE_EXPR:
+    if tag == TYPE_VAR_TUPLE_EXPR:
         return mypy.nodes.TypeVarTupleExpr.read(data)
-    assert False, f"Unknown symbol marker {marker}"
+    assert False, f"Unknown symbol tag {tag}"
 
 
 def read_overload_part(data: Buffer) -> OverloadPart:
-    marker = read_int(data)
-    if marker == DECORATOR:
+    tag = read_tag(data)
+    if tag == DECORATOR:
         return Decorator.read(data)
-    if marker == FUNC_DEF:
+    if tag == FUNC_DEF:
         return FuncDef.read(data)
-    assert False, f"Invalid marker for an OverloadPart {marker}"
+    assert False, f"Invalid tag for an OverloadPart {tag}"