1 files changed, 41 insertions, 0 deletions
diff --git a/scripts/zstd.py b/scripts/zstd.py
new file mode 100644
index 0000000..f85e7da
--- /dev/null
+++ b/scripts/zstd.py
@@ -0,0 +1,41 @@
+import subprocess
+from rcc_format import *
+
+def zstd(filelist, args):
+    if args.train is not None:
+        train(filelist, args.train, args.binary, args.dsize)
+        return
+
+    write_header(args.output, args.namespace)
+
+    for f in filelist:
+        with open(f.path, 'rb') as contents:
+            write_item(args.output, f.variable, compress(contents, args.binary, args.level, dictionary=args.dict))
+    
+    write_entries(args.output, filelist)
+    if args.dict is not None:
+        write_item(args.output, 'dict', args.dict.read())
+        print("constexpr auto dictionary = std::span(dict, dict_len);", file=args.output)
+    
+    print("constexpr auto compression = embed::Zstd;", file=args.output)
+    
+    write_footer(args.output, args.namespace)
+
+def train(filelist, output, zstd_bin, maxdict):
+    cmd = [ zstd_bin, '--train', '--maxdict=' + str(maxdict), '-o', output.name ]
+
+    for f in filelist:
+        cmd.append(f.path)
+
+    subprocess.run(cmd)
+
+def compress(file, zstd_bin, level, dictionary=None):
+    cmd = [ zstd_bin, '--compress', '--stdout', '-' + str(level) ]
+
+    if dictionary is not None:
+        cmd.append('-D')
+        cmd.append(dictionary.name)
+    
+    cmd.append(file.name)
+    return subprocess.run(cmd, capture_output=True).stdout
+