Change rcc command line

rcc:
    - add - and Zstd modes
    - fix Zstd dictionary train

4 files changed, 143 insertions, 91 deletions

diff --git a/meson.build b/meson.build
index f0f8194..3cc04f1 100644
--- a/meson.build
+++ b/meson.build
@@ -29,12 +29,19 @@ prog_python = import('python').find_installation('python3')
 resources_h = custom_target('resources.h',
   output: 'resources.h',
   input: 'scripts/rcc',
-  command: [prog_python, '@INPUT@', '--namespace=staticdata', '--output=@OUTPUT@', files('test/resources.xrc')],
+  command: [prog_python, '@INPUT@', '--namespace=staticdata', '--output=@OUTPUT@', '-', files('test/resources.xrc')],
 )
+
+zstd_dictionary = custom_target('zstd_dictionary',
+  output: 'zstd_dict',
+  input: 'scripts/rcc',
+  command: [ prog_python, '@INPUT@', 'Zstd', '--train=@OUTPUT@', files('test/resources.xrc') ]
+)
+
 zstd_resources_h = custom_target('zstd_resources.h',
   output: 'zstd_resources.h',
-  input: 'scripts/rcc',
-  command: [prog_python, '@INPUT@', '--namespace=zstd_data', '--compress=Zstd', '--dict=dictionary', '--output=@OUTPUT@', files('test/resources.xrc')],
+  input: [ 'scripts/rcc', zstd_dictionary ],
+  command: [prog_python, '@INPUT0@', '--namespace=zstd_data', '--output=@OUTPUT@', 'Zstd', '--dict=@INPUT1@', files('test/resources.xrc')],
 )
 
 test('libembed',
diff --git a/scripts/rcc b/scripts/rcc
index b81cf6e..479b609 100755
--- a/scripts/rcc
+++ b/scripts/rcc
@@ -2,106 +2,48 @@
 
 import argparse
 import sys
-import os.path
-import subprocess
-import xml.etree.ElementTree as xml
+from zstd import zstd
+from rcc_format import *
 
-def train(files, output, zstd='zstd', maxdict=512):
-    cmd = [ zstd, '--train', '--stdout', '--maxdict=' + str(maxdict), '-o', output ]
+def none(filelist, args):
+    write_header(args.output, args.namespace)
 
-    for f in files:
-        cmd.append(f.name)
-
-    subprocess.run(cmd)
-
-def compress(file, zstd='zstd', level=19, dictionary=None):
-    cmd = [ zstd, '--compress', '--stdout', '-' + str(level) ]
-
-    if dictionary is not None:
-        cmd.append('-D')
-        cmd.append(dictionary.name)
+    for f in filelist:
+        with open(f.path, 'rb') as contents:
+            write_item(args.output, f.variable, contents.read())
     
-    cmd.append(file.name)
-    return subprocess.run(cmd, capture_output=True).stdout
-
-def hexdump(array_name, array_data, out_h):
-    array_len = 0
+    write_entries(args.output, filelist)
+    print("constexpr auto compression = embed::None;", file=args.output)
     
-    print("constexpr unsigned char {}[] = {{".format(array_name), file=out_h)
-
-    for byte in array_data[0:len(array_data)]:
-        array_len+=1
-        if array_len%16 == 0:
-            print(" 0x{:02X},".format(byte), file=out_h)
-        else:
-            print(" 0x{:02X},".format(byte), file=out_h, end='')
-
-
-    print("};", file=out_h)
-    print("constexpr size_t {}_len = {};\n".format(array_name, array_len), file=out_h)
-
-def name(path):
-    name = path.replace('/', '_')
-    if name.endswith('.zstd'):
-        name = name[:-5]
-    name = name.replace('-', '_')
-    name = name.replace('.', '_')
-    return name
+    write_footer(args.output, args.namespace)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description='Resource Compiler for C++',
-        epilog='If using compression, make sure the required dependencies are provided.',
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+        epilog='For a full list of compression options, check {mode} --help.',
     )
 
-    parser.add_argument('input', type=argparse.FileType('rt'), help='input file (.xrc)')
-    parser.add_argument('-o', '--output', type=argparse.FileType('wt'), metavar='OUT', default=sys.stdout, help='output header file')
-    parser.add_argument('-n', '--namespace', type=str, help='namespace')
-
-    parser.add_argument('--compress', choices=[ 'None', 'Zstd' ], default='None', help='compress input files using algorightm')
-    parser.add_argument('--dict', type=argparse.FileType('rb'), help='[zstd] use specified dictionary, recommended for many similar small files')
-    parser.add_argument('--train', action='store_true', help='[zstd] train dictionary')
+    mode = parser.add_subparsers(help='compression mode')
 
-    args=parser.parse_args()
-
-    entries_list = ""
-
-    if args.train:
-        train(args.input, args.dict.name)
-
-    # write header
-    print("// Autogenerated binary file hexdump", file=args.output)
-    print("// This file may get overwritten by the build system\n", file=args.output)
-    print("#include <embed.h>\n", file=args.output)
-    print("namespace {} {{".format(args.namespace), file=args.output)
-
-    # write file data
-    for child in xml.parse(args.input).getroot():
-        if child.tag == 'qresource':
-            prefix = child.attrib['prefix']
-            for i in child:
-                vname = name(i.text)
-                with open(i.text, 'rb') as f:
-                    if args.compress == 'None':
-                        hexdump(vname, f.read(), args.output)
-                    elif args.compress == 'Zstd':
-                        hexdump(vname, compress(f, dictionary=args.dict), args.output)
-                entries_list += "  {{ \"{}/{}\", std::span({}, {}_len) }},\n".format(prefix, i.attrib['alias'], vname, vname)
-
-    # write dictionary
-    if args.dict is not None:
-        hexdump('dict', args.dict.read(), args.output)
+    none_mode = mode.add_parser('-')
+    none_mode.set_defaults(func=none)
     
-    # write entries
-    print("constexpr auto entries = frozen::make_unordered_map<frozen::string, std::span<const unsigned char>>({", file=args.output)
-    print(entries_list, file=args.output)
-    print("});\n", file=args.output)
+    zstd_mode = mode.add_parser('Zstd',
+        description='use Zstd compression',
+        epilog='A dictionary is recommended if compressing many small files. size(source)/size(dictionary) should be >= 10'
+    )
+    zstd_mode.add_argument('--binary', type=str, default='zstd', help='zstd binary name')
+    zstd_mode.add_argument('--train', type=argparse.FileType('wb'), help='train dictionary and exit')
+    zstd_mode.add_argument('-d', '--dict', type=argparse.FileType('rb'), help='use dictionary, recommended for many similar small files')
+    zstd_mode.add_argument('--dsize', type=int, default=512, help='dictionary size, used for training')
+    zstd_mode.add_argument('-l', '--level', type=int, default=19, help='compression level')
+    zstd_mode.set_defaults(func=zstd)
+
+    parser.add_argument('input', type=argparse.FileType('rt'), help='input file (.xrc)')
+    parser.add_argument('-o', '--output', type=argparse.FileType('wt'), default=sys.stdout, help='output header file')
+    parser.add_argument('-n', '--namespace', type=str, default='resources', help='namespace')
     
-    # write metadata
-    print("constexpr auto compression = embed::{};".format(args.compress), file=args.output)
-    if args.dict is not None:
-        print("constexpr auto dictionary = std::span(dict, dict_len);", file=args.output)
+    args=parser.parse_args()
 
-    print("}} // namespace {}".format(args.namespace), file=args.output)
+    args.func(filelist(args.input), args)
 
diff --git a/scripts/rcc_format.py b/scripts/rcc_format.py
new file mode 100644
index 0000000..d0ab06d
--- /dev/null
+++ b/scripts/rcc_format.py
@@ -0,0 +1,62 @@
+from collections import namedtuple
+import xml.etree.ElementTree as xml
+
+resource = namedtuple('resource', 'alias variable path')
+
+def to_variable_name(path):
+    name = path.replace('/', '_')
+    if name.endswith('.zstd'):
+        name = name[:-5]
+    name = name.replace('-', '_')
+    name = name.replace('.', '_')
+    return name
+
+def filelist(file):
+    root = xml.parse(file).getroot()
+    if root.tag != 'RCC':
+        return None
+
+    files = []
+    for child in root:
+        if child.tag == 'qresource':
+            prefix = child.attrib['prefix']
+            for i in child:
+                alias = prefix + '/' + i.attrib['alias']
+                variable = to_variable_name(i.text)
+                path = i.text
+                files.append(resource(alias, variable, path))
+
+    return files
+
+def write_header(file, namespace):
+    print("// Autogenerated binary file hexdump", file=file)
+    print("// This file may get overwritten by the build system\n", file=file)
+    print("#include <embed.h>\n", file=file)
+    print("namespace {} {{".format(namespace), file=file)
+
+def write_item(file, array_name, array_data):
+    line_items = 0
+    
+    print("constexpr unsigned char {}[] = {{".format(array_name), file=file)
+
+    for byte in array_data[0:len(array_data)]:
+        line_items+=1
+        if line_items == 16:
+            print(" 0x{:02X},".format(byte), file=file)
+            line_items = 0
+        else:
+            print(" 0x{:02X},".format(byte), file=file, end='')
+
+
+    print("};", file=file)
+    print("constexpr size_t {}_len = {};\n".format(array_name, len(array_data)), file=file)
+
+def write_entries(file, resource_list):
+    print("constexpr auto entries = frozen::make_unordered_map<frozen::string, std::span<const unsigned char>>({", file=file)
+    for f in resource_list:
+        print("  {{ \"{}\", std::span({}, {}_len) }},".format(f.alias, f.variable, f.variable), file=file)
+    print("});\n", file=file)
+
+def write_footer(file, namespace):
+    print("\n}} // namespace {}".format(namespace), file=file)
+
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
+