feat(cli): generate static files at the granularity of proto messages (…

…protobufjs#1840) * feat: add message filter for cli * feat: add test * fix: update comment * fix: update error message * fix: remove test file * fix: lint, jsdoc comments, return values Co-authored-by: pluschen <pluschen@tencent.com> Co-authored-by: Alexander Fenster <fenster@google.com>
Mu-L · Jan 24, 2023 · 32f2d6a · 32f2d6a
1 parent ea7b9a6
commit 32f2d6a
Show file tree

Hide file tree

Showing 7 changed files with 225 additions and 1 deletion.
diff --git a/cli/pbjs.js b/cli/pbjs.js
@@ -40,7 +40,7 @@ exports.main = function main(args, callback) {
             "force-long": "strict-long",
             "force-message": "strict-message"
         },
-        string: [ "target", "out", "path", "wrap", "dependency", "root", "lint" ],
+        string: [ "target", "out", "path", "wrap", "dependency", "root", "lint", "filter" ],
         boolean: [ "create", "encode", "decode", "verify", "convert", "delimited", "typeurl", "beautify", "comments", "service", "es6", "sparse", "keep-case", "alt-comment", "force-long", "force-number", "force-enum-string", "force-message", "null-defaults" ],
         default: {
             target: "json",
@@ -98,6 +98,9 @@ exports.main = function main(args, callback) {
                 "",
                 "  -p, --path       Adds a directory to the include path.",
                 "",
+                "  --filter         Set up a filter to configure only those messages you need and their dependencies to compile, this will effectively reduce the final file size",
+                "                   Set A json file path, Example of file content: {\"messageNames\":[\"mypackage.messageName1\", \"messageName2\"] } ",
+                "",
                 "  -o, --out        Saves to a file instead of writing to stdout.",
                 "",
                 "  --sparse         Exports only those types referenced from a main file (experimental).",
@@ -308,7 +311,20 @@ exports.main = function main(args, callback) {
         root.resolveAll();
     }
 
+    function filterMessage() {
+        if (argv.filter) {
+            // This is a piece of degradable logic
+            try {
+                const needMessage = JSON.parse(fs.readFileSync(argv.filter));
+                util.filterMessage(root, needMessage);
+            } catch (error) {
+                process.stderr.write(`The filter not work, please check whether the file is correct: ${error.message}\n`);
+            }
+        }
+    }
+
     function callTarget() {
+        filterMessage();
         target(root, argv, function targetCallback(err, output) {
             if (err) {
                 if (callback)

diff --git a/cli/util.js b/cli/util.js
@@ -125,3 +125,119 @@ exports.pad = function(str, len, l) {
     return str;
 };
 
+
+/**
+ * DFS to get all message dependencies, cache in filterMap.
+ * @param {Root} root  The protobuf root instance
+ * @param {Message} message  The message need to process.
+ * @param {Map} filterMap  The result of message you need and their dependencies.
+ * @param {Map} flatMap  A flag to record whether the message was searched.
+ * @returns {undefined}  Does not return a value
+ */
+function dfsFilterMessageDependencies(root, message, filterMap, flatMap) {
+    if (message instanceof protobuf.Type) {
+        if (flatMap.get(`${message.fullName}`)) return;
+        flatMap.set(`${message.fullName}`, true);
+        for (var field of message.fieldsArray) {
+            if (field.resolvedType) {
+                // a nested message
+                if (field.resolvedType.parent.name === message.name) {
+                    var nestedMessage = message.nested[field.resolvedType.name];
+                    dfsFilterMessageDependencies(root, nestedMessage, filterMap, flatMap);
+                    continue;
+                }
+                var packageName = field.resolvedType.parent.name;
+                var typeName = field.resolvedType.name;
+                var fullName = packageName ? `${packageName}.${typeName}` : typeName;
+                doFilterMessage(root, { messageNames: [fullName] }, filterMap, flatMap, packageName);
+            }
+        }
+    }
+}
+
+/**
+ * DFS to get all message you need and their dependencies, cache in filterMap.
+ * @param {Root} root  The protobuf root instance
+ * @param {object} needMessageConfig  Need message config:
+ * @param {string[]} needMessageConfig.messageNames  The message names array in the root namespace you need to gen. example: [msg1, msg2]
+ * @param {Map} filterMap The result of message you need and their dependencies.
+ * @param {Map} flatMap A flag to record whether the message was searched.
+ * @param {string} currentPackageName  Current package name
+ * @returns {undefined}  Does not return a value
+ */
+function doFilterMessage(root, needMessageConfig, filterMap, flatMap, currentPackageName) {
+    var needMessageNames = needMessageConfig.messageNames;
+
+    for (var messageFullName of needMessageNames) {
+        var nameSplit = messageFullName.split(".");
+        var packageName = "";
+        var messageName = "";
+        if (nameSplit.length > 1) {
+            packageName = nameSplit[0];
+            messageName = nameSplit[1];
+        } else {
+            messageName = nameSplit[0];
+        }
+
+        // in Namespace
+        if (packageName) {
+            var ns = root.nested[packageName];
+            if (!ns || !(ns instanceof protobuf.Namespace)) {
+                throw new Error(`package not foud ${currentPackageName}.${messageName}`);
+            }
+
+            doFilterMessage(root, { messageNames: [messageName] }, filterMap, flatMap, packageName);
+        } else {
+            var message = root.nested[messageName];
+
+            if (currentPackageName) {
+                message = root.nested[currentPackageName].nested[messageName];
+            }
+
+            if (!message) {
+                throw new Error(`message not foud ${currentPackageName}.${messageName}`);
+            }
+
+            var set = filterMap.get(currentPackageName);
+            if (!filterMap.has(currentPackageName)) {
+                set = new Set();
+                filterMap.set(currentPackageName, set);
+            }
+
+            set.add(messageName);
+
+            // dfs to find all dependencies
+            dfsFilterMessageDependencies(root, message, filterMap, flatMap, currentPackageName);
+        }
+    }
+}
+
+/**
+ * filter the message you need and their dependencies, all others will be delete from root.
+ * @param {Root} root  Root the protobuf root instance
+ * @param {object} needMessageConfig  Need message config:
+ * @param {string[]} needMessageConfig.messageNames  Tthe message names array in the root namespace you need to gen. example: [msg1, msg2]
+ * @returns {boolean} True if a message should present in the generated files
+ */
+exports.filterMessage = function (root, needMessageConfig) {
+    var filterMap = new Map();
+    var flatMap = new Map();
+    doFilterMessage(root, needMessageConfig, filterMap, flatMap, "");
+    root._nestedArray = root._nestedArray.filter(ns => {
+        if (ns instanceof protobuf.Type || ns instanceof protobuf.Enum) {
+            return filterMap.get("").has(ns.name);
+        } else if (ns instanceof protobuf.Namespace) {
+            if (!filterMap.has(ns.name)) {
+                return false;
+            }
+            ns._nestedArray = ns._nestedArray.filter(nns => {
+                const nnsSet = filterMap.get(ns.name);
+                return nnsSet.has(nns.name);
+            });
+
+            return true;
+        }
+        return true;
+    });
+};
+
diff --git a/package.json b/package.json
@@ -11,6 +11,14 @@
   "engines": {
     "node": ">=12.0.0"
   },
+  "eslintConfig": {
+    "env": {
+      "es6": true
+    },
+    "parserOptions": {
+      "ecmaVersion": 6
+    }
+  },
   "keywords": [
     "protobuf",
     "protocol-buffers",

diff --git a/tests/cli.js b/tests/cli.js
@@ -4,6 +4,7 @@ var tape = require("tape");
 var path = require("path");
 var Module = require("module");
 var protobuf = require("..");
+var fs = require("fs");
 
 function cliTest(test, testFunc) {
     // pbjs does not seem to work with Node v4, so skip this test if we're running on it
@@ -162,3 +163,54 @@ tape.test("with null-defaults, absent optional fields have null values", functio
         });
     });
 });
+
+
+tape.test("pbjs generates static code with message filter", function (test) {
+    cliTest(test, function () {
+        var root = protobuf.loadSync("tests/data/cli/test-filter.proto");
+        root.resolveAll();
+
+        var staticTarget = require("../cli/targets/static");
+        var util = require("../cli/util");
+
+        const needMessageConfig = JSON.parse(fs.readFileSync("tests/data/cli/filter.json"));
+
+        util.filterMessage(root, needMessageConfig);
+
+        staticTarget(root, {
+            create: true,
+            decode: true,
+            encode: true,
+            convert: true,
+            "null-defaults": true,
+        }, function (err, jsCode) {
+            test.error(err, 'static code generation worked');
+
+            // jsCode is the generated code; we'll eval it
+            // (since this is what we normally does with the code, right?)
+            // This is a test code. Do not use this in production.
+            var $protobuf = protobuf;
+            eval(jsCode);
+
+            console.log(protobuf.roots);
+
+            var NeedMessage1 = protobuf.roots.default.filtertest.NeedMessage1;
+            var NeedMessage2 = protobuf.roots.default.filtertest.NeedMessage2;
+            var DependentMessage1 = protobuf.roots.default.filtertest.DependentMessage1;
+            var DependentMessageFromImport = protobuf.roots.default.DependentMessageFromImport;
+
+            var NotNeedMessageInRootFile = protobuf.roots.default.filtertest.NotNeedMessageInRootFile;
+            var NotNeedMessageInImportFile = protobuf.roots.default.NotNeedMessageInImportFile;
+
+            test.ok(NeedMessage1, "NeedMessage1 is loaded");
+            test.ok(NeedMessage2, "NeedMessage2 is loaded");
+            test.ok(DependentMessage1, "DependentMessage1 is loaded");
+            test.ok(DependentMessageFromImport, "DependentMessageFromImport is loaded");
+
+            test.notOk(NotNeedMessageInImportFile, "NotNeedMessageInImportFile is not loaded");
+            test.notOk(NotNeedMessageInRootFile, "NotNeedMessageInRootFile is not loaded");
+
+            test.end();
+        });
+    });
+});
diff --git a/tests/data/cli/filter.json b/tests/data/cli/filter.json
@@ -0,0 +1,3 @@
+{
+  "messageNames": ["filtertest.NeedMessage1", "filtertest.NeedMessage2"]
+}
diff --git a/tests/data/cli/test-filter-import.proto b/tests/data/cli/test-filter-import.proto
@@ -0,0 +1,8 @@
+
+message DependentMessageFromImport {
+  optional int32 test1 = 1;
+}
+
+message NotNeedMessageInImportFile {
+  optional int32 test1 = 1;
+}
diff --git a/tests/data/cli/test-filter.proto b/tests/data/cli/test-filter.proto
@@ -0,0 +1,21 @@
+package filtertest;
+import "./test-filter-import.proto";
+
+message NeedMessage1 {
+  optional uint32 test1 = 1;
+  optional NeedMessage2 needMessage2 = 2;
+  optional DependentMessage1 dependentMessage1 = 3;
+  optional DependentMessageFromImport dependentMessage2 = 4;
+}
+
+message NeedMessage2 {
+  optional uint32 test1 = 1;
+}
+
+message DependentMessage1 {
+  optional uint32 test1 = 1;
+}
+
+message NotNeedMessageInRootFile {
+  optional uint32 test1 = 1;
+}