fix(bake): use cross-platform Zstandard for archives and fixtures

nx2/bake_ds/fixture.py

@@ -249,12 +249,28 @@ def write(rel, data):
         with open(path, "wb") as f:
             f.write(data)
 
-    def z(data):
+    def z(data, dictionary=None):
         try:
-            from compression import zstd   # (Python 3.14)
-            return zstd.compress(data)
-        except Exception:
-            return subprocess.run(["zstd", "-q", "-c"], input=data, stdout=subprocess.PIPE, check=True).stdout
+            import zstandard
+        except ImportError:
+            if dictionary is None:
+                try:
+                    from compression import zstd   # Python 3.14
+                except ImportError:
+                    pass
+                else:
+                    return zstd.compress(data)
+            import tempfile
+            with tempfile.TemporaryDirectory(prefix="bake-dictionary-") as directory:
+                cmd = ["zstd", "-q", "-c"]
+                if dictionary is not None:
+                    path = os.path.join(directory, "dictionary")
+                    with open(path, "wb") as f:
+                        f.write(dictionary)
+                    cmd += ["-D", path]
+                return subprocess.run(cmd, input=data, stdout=subprocess.PIPE, check=True).stdout
+        decoded = None if dictionary is None else zstandard.ZstdCompressionDict(dictionary)
+        return zstandard.ZstdCompressor(dict_data=decoded).compress(data)
 
     write("Shader/Test.bfsha.zs", z(bnsh([graphics_pair(vid=1), compute(local=(8, 4, 1))])))
     inner_bnsh = bnsh([graphics_pair(vid=2)])
@@ -273,21 +289,9 @@ def z(data):
         # (a raw dictionary; libzstd accepts raw content dictionaries, which is what the pack holds)
         d = (b"bake-ds-dictionary " * 256)[:4096]
         write("Pack/ZsDic.pack.zs", z(dic_pack([d])))
-        import tempfile
-        with tempfile.NamedTemporaryFile(delete=False) as tf:
-            tf.write(d)
-            path = tf.name
-        try:
-            data = bnsh([graphics_pair(vid=10)])
-            comp = subprocess.run(["zstd", "-q", "-c", "-D", path], input=data, stdout=subprocess.PIPE,
-                                  stderr=subprocess.PIPE)
-            if comp.returncode == 0 and comp.stdout[:4] == b"\x28\xb5\x2f\xfd":
-                write("Shader/Dict.bfsha.zs", comp.stdout)
-                stats["dict"] = "written"
-            else:
-                stats["dict"] = "zstd -D failed"
-        finally:
-            os.unlink(path)
+        data = bnsh([graphics_pair(vid=10)])
+        write("Shader/Dict.bfsha.zs", z(data, d))
+        stats["dict"] = "written"
     return stats
 
 

nx2/bake_ds/romenum.py

@@ -4,7 +4,7 @@
 
 The walk and the containers are the ones nx2/shader_census.py documents (BNSH variations, agl's BAHS,
 NDSR networks, NDVE upscaler tables, Effect's .vsetb BNSH, SARC wrappers, zstd with the game's
-dictionaries). This module is standalone: no zstandard package, no game files in the repository.
+dictionaries). This module prefers the project zstandard package, with a system libzstd fallback. No game files are in the repository.
 The only thing it shares with the census is the format, not code or files.
 
 Output: the same HKSEED01 the back end reads (gpu/vk/vk_shader.cpp:1844 seed_open):
@@ -19,9 +19,8 @@
 so either reads, but FNV dedupes exactly what the runtime treats as one shader. (Bake-Mimo made the same
 point for its stable IDs.)
 
-zstd: libzstd through ctypes (frames with and without a dictionary), because the game's archives are
-.zs and the dictionaries are player data. A plain zstd CLI fallback exists; nothing is written outside
---out/--manifest.
+zstd: the project's cross-platform zstandard package, or libzstd through ctypes when unavailable.
+Frames and dictionaries remain player data. Nothing is written outside --out/--manifest.
 """
 import argparse
 import collections
@@ -49,8 +48,17 @@ class _Zstd:
     ERROR = (1 << 64) - 2
 
     def __init__(self):
-        name = ctypes.util.find_library("zstd") or "libzstd.so.1"
-        self.lib = ctypes.CDLL(name)
+        candidates = [ctypes.util.find_library("zstd")]
+        candidates += (["libzstd.dll", "zstd.dll"] if os.name == "nt" else
+                       ["libzstd.so.1", "libzstd.so", "libzstd.dylib"])
+        for name in filter(None, candidates):
+            try:
+                self.lib = ctypes.CDLL(name)
+                break
+            except OSError:
+                continue
+        else:
+            raise RuntimeError("zstd unavailable: install the Python zstandard package or system libzstd")
         L = self.lib
         L.ZSTD_createDCtx.restype = ctypes.c_void_p
         L.ZSTD_freeDCtx.argtypes = [ctypes.c_void_p]
@@ -102,13 +110,40 @@ def decompress(self, data, dictionaries=()):
             cap *= 4   # (a frame with no content size: grow, as a stream would)
 
 
+class _PythonZstd:
+    def __init__(self, module):
+        self.module = module
+
+    def decompress(self, data, dictionaries=()):
+        if data[:4] != ZSTD_MAGIC:
+            return data
+        last = None
+        for dictionary in (None, *dictionaries):
+            try:
+                decoded_dictionary = (None if dictionary is None else
+                                      self.module.ZstdCompressionDict(dictionary))
+                decoder = self.module.ZstdDecompressor(dict_data=decoded_dictionary).decompressobj()
+                result = decoder.decompress(data)
+                if not decoder.eof:
+                    raise ValueError("incomplete zstd frame")
+                return result
+            except (self.module.ZstdError, ValueError) as exc:
+                last = exc
+        raise ValueError("zstd: %s" % last)
+
+
 _ZSTD = None
 
 
 def zstd():
     global _ZSTD
     if _ZSTD is None:
-        _ZSTD = _Zstd()
+        try:
+            import zstandard
+        except ImportError:
+            _ZSTD = _Zstd()
+        else:
+            _ZSTD = _PythonZstd(zstandard)
     return _ZSTD
 
 

tests/test_bake_ds.py

@@ -19,6 +19,8 @@
 import sys
 import tempfile
 
+RUN_ALL_MODULES = ["zstandard"]
+
 ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.insert(0, os.path.join(ROOT, "nx2"))
 sys.path.insert(0, os.path.join(ROOT, "nx2", "bake_ds"))

tests/test_bake_zstd.py

@@ -0,0 +1,64 @@
+"""Synthetic Zstandard frames work without a platform-specific shared library."""
+from pathlib import Path
+import sys
+import unittest
+from unittest import mock
+
+RUN_ALL_MODULES = ['zstandard']
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'nx2'))
+import zstandard as zs
+from bake_ds import romenum
+
+
+class PortableZstd(unittest.TestCase):
+    def setUp(self):
+        self.decoder = romenum._PythonZstd(zs)
+        self.payload = b'Longshot synthetic compression regression\n' * 8192
+
+    def test_known_and_unknown_size(self):
+        for sized in (False, True):
+            with self.subTest(sized=sized):
+                frame = zs.ZstdCompressor(write_content_size=sized).compress(self.payload)
+                self.assertEqual(self.decoder.decompress(frame), self.payload)
+
+    def test_required_trained_dictionary_and_fallback_order(self):
+        samples = [(b'synthetic dictionary sample %d ' % i) * 64 for i in range(128)]
+        dictionary = zs.train_dictionary(2048, samples)
+        frame = zs.ZstdCompressor(dict_data=dictionary).compress(samples[31])
+        with self.assertRaises(ValueError):
+            self.decoder.decompress(frame)
+        self.assertEqual(self.decoder.decompress(frame, [b'wrong dictionary', dictionary.as_bytes()]),
+                         samples[31])
+
+    def test_raw_dictionary(self):
+        raw = self.payload[:4096]
+        frame = zs.ZstdCompressor(dict_data=zs.ZstdCompressionDict(raw)).compress(self.payload)
+        self.assertEqual(self.decoder.decompress(frame, [raw]), self.payload)
+
+    def test_truncated_and_invalid_frames_are_errors(self):
+        for sized in (False, True):
+            frame = zs.ZstdCompressor(write_content_size=sized).compress(self.payload)
+            with self.assertRaises(ValueError):
+                self.decoder.decompress(frame[:-1])
+        with self.assertRaises(ValueError):
+            self.decoder.decompress(romenum.ZSTD_MAGIC + b'invalid')
+
+    def test_plain_input(self):
+        self.assertEqual(self.decoder.decompress(b'plain'), b'plain')
+
+    def test_python_backend_never_loads_shared_library(self):
+        with mock.patch.object(romenum, '_ZSTD', None), \
+             mock.patch.object(romenum.ctypes, 'CDLL', side_effect=AssertionError('DLL load')):
+            self.assertIsInstance(romenum.zstd(), romenum._PythonZstd)
+            self.assertEqual(romenum.zstd().decompress(zs.ZstdCompressor().compress(self.payload)),
+                             self.payload)
+
+    def test_missing_native_library_has_actionable_error(self):
+        with mock.patch.object(romenum.ctypes.util, 'find_library', return_value=None), \
+             mock.patch.object(romenum.ctypes, 'CDLL', side_effect=OSError('missing')):
+            with self.assertRaisesRegex(RuntimeError, 'install the Python zstandard package'):
+                romenum._Zstd()
+
+
+if __name__ == '__main__':
+    unittest.main()