Minor changes

Build/lib/misc.ts

@@ -3,6 +3,8 @@ import fs from 'node:fs';
 import type { PathLike } from 'node:fs';
 import fsp from 'node:fs/promises';
 
+export type MaybePromise<T> = T | Promise<T>;
+
 export function fastStringCompare(a: string, b: string) {
   const lenA = a.length;
   const lenB = b.length;

Build/lib/rules/base.ts

@@ -4,6 +4,7 @@ import { invariant, not } from 'foxts/guard';
 import picocolors from 'picocolors';
 import fs from 'node:fs';
 import { writeFile } from '../misc';
+import type { MaybePromise } from '../misc';
 import { fastStringArrayJoin } from 'foxts/fast-string-array-join';
 import { readFileByLine } from '../fetch-text-by-line';
 import { asyncWriteToStream } from 'foxts/async-write-to-stream';
@@ -13,6 +14,10 @@ import { createRetrieKeywordFilter as createKeywordFilter } from 'foxts/retrie';
 import path from 'node:path';
 import { SurgeMitmSgmodule } from '../writing-strategy/surge';
 
+/**
+ * Holds the universal rule data (domain, ip, url-regex, etc. etc.)
+ * This class is not about format, instead it will call the class that does
+ */
 export class FileOutput {
   protected strategies: Array<BaseWriteStrategy | false> = [];
 
@@ -135,12 +140,24 @@ export class FileOutput {
     }
   }
 
-  addFromDomainset(source: AsyncIterable<string> | Iterable<string> | string[]) {
-    this.pendingPromise = (this.pendingPromise ||= Promise.resolve()).then(() => this.addFromDomainsetPromise(source));
+  addFromDomainset(source: MaybePromise<AsyncIterable<string> | Iterable<string> | string[]>) {
+    if (this.pendingPromise) {
+      if ('then' in source) {
+        this.pendingPromise = this.pendingPromise.then(() => source).then(src => this.addFromDomainsetPromise(src));
+        return this;
+      }
+      this.pendingPromise = this.pendingPromise.then(() => this.addFromDomainsetPromise(source));
+      return this;
+    }
+    if ('then' in source) {
+      this.pendingPromise = source.then(src => this.addFromDomainsetPromise(src));
+      return this;
+    }
+    this.pendingPromise = this.addFromDomainsetPromise(source);
     return this;
   }
 
-  private async addFromRulesetPromise(source: AsyncIterable<string> | Iterable<string>) {
+  private async addFromRulesetPromise(source: AsyncIterable<string> | Iterable<string> | string[]) {
     for await (const line of source) {
       const splitted = line.split(',');
       const type = splitted[0];
@@ -203,13 +220,20 @@ export class FileOutput {
     }
   }
 
-  addFromRuleset(source: AsyncIterable<string> | Iterable<string> | Promise<Iterable<string>>) {
+  addFromRuleset(source: MaybePromise<AsyncIterable<string> | Iterable<string>>) {
     if (this.pendingPromise) {
-      this.pendingPromise = this.pendingPromise.then(() => source);
-    } else {
-      this.pendingPromise = Promise.resolve(source);
+      if ('then' in source) {
+        this.pendingPromise = this.pendingPromise.then(() => source).then(src => this.addFromRulesetPromise(src));
+        return this;
+      }
+      this.pendingPromise = this.pendingPromise.then(() => this.addFromRulesetPromise(source));
+      return this;
     }
-    this.pendingPromise = this.pendingPromise.then((source) => this.addFromRulesetPromise(source));
+    if ('then' in source) {
+      this.pendingPromise = source.then(src => this.addFromRulesetPromise(src));
+      return this;
+    }
+    this.pendingPromise = this.addFromRulesetPromise(source);
     return this;
   }
 
@@ -282,15 +306,16 @@ export class FileOutput {
   // }
   private strategiesWritten = false;
 
-  private async writeToStrategies() {
+  private writeToStrategies() {
+    if (this.pendingPromise) {
+      throw new Error('You should call done() before calling writeToStrategies()');
+    }
     if (this.strategiesWritten) {
       throw new Error('Strategies already written');
     }
 
     this.strategiesWritten = true;
 
-    await this.done();
-
     const kwfilter = createKeywordFilter(Array.from(this.domainKeywords));
 
     if (this.strategies.filter(not(false)).length === 0) {
@@ -415,40 +440,44 @@ export class FileOutput {
     }
   }
 
-  write(): Promise<void> {
+  write(): Promise<unknown> {
     return this.span.traceChildAsync('write all', async (childSpan) => {
-      const promises: Array<Promise<void> | void> = [];
+      await this.done();
+      childSpan.traceChildSync('write to strategies', this.writeToStrategies.bind(this));
 
-      await childSpan.traceChildAsync('write to strategies', this.writeToStrategies.bind(this));
+      return childSpan.traceChildAsync('output to disk', (childSpan) => {
+        const promises: Array<Promise<void> | void> = [];
 
-      invariant(this.title, 'Missing title');
-      invariant(this.description, 'Missing description');
+        invariant(this.title, 'Missing title');
+        invariant(this.description, 'Missing description');
 
-      for (let i = 0, len = this.strategies.length; i < len; i++) {
-        const strategy = this.strategies[i];
-        if (strategy) {
-          const basename = (strategy.overwriteFilename || this.id) + '.' + strategy.fileExtension;
-          promises.push(strategy.output(
-            childSpan,
-            this.title,
-            this.description,
-            this.date,
-            path.join(
-              strategy.outputDir,
-              strategy.type
-                ? path.join(strategy.type, basename)
-                : basename
-            )
-          ));
+        for (let i = 0, len = this.strategies.length; i < len; i++) {
+          const strategy = this.strategies[i];
+          if (strategy) {
+            const basename = (strategy.overwriteFilename || this.id) + '.' + strategy.fileExtension;
+            promises.push(strategy.output(
+              childSpan,
+              this.title,
+              this.description,
+              this.date,
+              path.join(
+                strategy.outputDir,
+                strategy.type
+                  ? path.join(strategy.type, basename)
+                  : basename
+              )
+            ));
+          }
         }
-      }
 
-      await Promise.all(promises);
+        return Promise.all(promises);
+      });
     });
   }
 
   async compile(): Promise<Array<string[] | null>> {
-    await this.writeToStrategies();
+    await this.done();
+    this.writeToStrategies();
 
     return this.strategies.reduce<Array<string[] | null>>((acc, strategy) => {
       if (strategy) {

Build/lib/writing-strategy/base.ts

@@ -1,11 +1,24 @@
 import type { Span } from '../../trace';
 import { compareAndWriteFile } from '../create-file';
 
+/**
+ * The class is not about holding rule data, instead it determines how the
+ * date is written to a file.
+ */
 export abstract class BaseWriteStrategy {
+  /**
+   * Sometimes a ruleset will create extra files (e.g. reject-url-regex w/ mitm.sgmodule),
+   * and doesn't share the same filename and id. This property is used to overwrite the filename.
+   */
   public overwriteFilename: string | null = null;
+  public withFilename(filename: string) {
+    this.overwriteFilename = filename;
+    return this;
+  }
+
   public abstract readonly type: 'domainset' | 'non_ip' | 'ip' | (string & {});
 
-  abstract readonly fileExtension: 'conf' | 'txt' | 'json' | (string & {});
+  abstract readonly fileExtension: 'conf' | 'txt' | 'json' | 'sgmodule' /* | (string & {}) */;
 
   constructor(public readonly outputDir: string) {}
 
@@ -28,6 +41,8 @@ export abstract class BaseWriteStrategy {
   abstract writeDestinationPorts(port: Set<string>): void;
   abstract writeOtherRules(rule: string[]): void;
 
+  protected abstract withPadding(title: string, description: string[] | readonly string[], date: Date, content: string[]): string[];
+
   static readonly domainWildCardToRegex = (domain: string) => {
     let result = '^';
     for (let i = 0, len = domain.length; i < len; i++) {
@@ -49,8 +64,6 @@ export abstract class BaseWriteStrategy {
     return result;
   };
 
-  protected abstract withPadding(title: string, description: string[] | readonly string[], date: Date, content: string[]): string[];
-
   public output(
     span: Span,
     title: string,

Build/lib/writing-strategy/surge.ts

@@ -156,7 +156,7 @@ export class SurgeMitmSgmodule extends BaseWriteStrategy {
 
   constructor(moduleName: string, outputDir = OUTPUT_MODULES_DIR) {
     super(outputDir);
-    this.overwriteFilename = moduleName;
+    this.withFilename(moduleName);
   }
 
   writeDomain = noop;