optimize: trimmer config idl file preserve

tool/trimmer/trim/config.go

@@ -29,6 +29,7 @@ type YamlArguments struct {
 	Preserve         *bool    `yaml:"preserve,omitempty"`
 	PreservedStructs []string `yaml:"preserved_structs,omitempty"`
 	MatchGoName      *bool    `yaml:"match_go_name,omitempty"`
+	PreservedFiles   []string `yaml:"preserved_files,omitempty"`
 }
 
 func ParseYamlConfig(path string) *YamlArguments {

tool/trimmer/trim/mark.go

@@ -312,5 +312,27 @@ func (t *Trimmer) checkPreserve(theStruct *parser.StructLike) bool {
 			return true
 		}
 	}
-	return t.preserveRegex.MatchString(strings.ToLower(theStruct.ReservedComments))
+	if t.preserveRegex.MatchString(strings.ToLower(theStruct.ReservedComments)) {
+		return true
+	}
+	// 如果整个文件也是要保留的，那么里面的结构体也不删除
+	if t.preserveFileStructs[theStruct] {
+		return true
+	}
+	return false
+}
+
+func (t *Trimmer) loadPreserveFiles(ast *parser.Thrift, preserveFiles []string) {
+	preserveFilesMap := map[string]bool{}
+	for _, fn := range preserveFiles {
+		preserveFilesMap[fn] = true
+	}
+	t.preserveFileStructs = map[*parser.StructLike]bool{}
+	for th := range ast.DepthFirstSearch() {
+		if preserveFilesMap[th.Filename] {
+			for _, st := range ast.Structs {
+				t.preserveFileStructs[st] = true
+			}
+		}
+	}
 }

tool/trimmer/trim/trimmer.go

@@ -36,15 +36,17 @@ type Trimmer struct {
 	marks  map[string]map[interface{}]bool
 	outDir string
 	// use -m
-	trimMethods      []*regexp2.Regexp
-	matchGoName      bool
-	trimMethodValid  []bool
-	preserveRegex    *regexp.Regexp
-	forceTrimming    bool
-	preservedStructs []string
-	structsTrimmed   int
-	fieldsTrimmed    int
-	extServices      []*parser.Service
+	trimMethods         []*regexp2.Regexp
+	matchGoName         bool
+	trimMethodValid     []bool
+	preserveRegex       *regexp.Regexp
+	forceTrimming       bool
+	preservedStructs    []string
+	structsTrimmed      int
+	fieldsTrimmed       int
+	extServices         []*parser.Service
+	PreservedFiles      []string
+	preserveFileStructs map[*parser.StructLike]bool
 }
 
 type TrimASTArg struct {
@@ -53,6 +55,7 @@ type TrimASTArg struct {
 	Preserve        *bool
 	MatchGoName     *bool
 	PreserveStructs []string
+	PreservedFiles  []string
 }
 
 type TrimResultInfo struct {
@@ -80,8 +83,9 @@ func (t *TrimResultInfo) FieldTrimmedPercentage() float64 {
 
 // TrimAST parse the cfg and trim the single AST
 func TrimAST(arg *TrimASTArg) (trimResultInfo *TrimResultInfo, err error) {
-	var preservedStructs []string
+	var preservedStructs, preservedFiles []string
 	preservedStructs = arg.PreserveStructs
+	preservedFiles = arg.PreservedFiles
 	if wd, err := dir_utils.Getwd(); err == nil {
 		cfg := ParseYamlConfig(wd)
 		if cfg != nil {
@@ -98,6 +102,9 @@ func TrimAST(arg *TrimASTArg) (trimResultInfo *TrimResultInfo, err error) {
 			if len(preservedStructs) == 0 {
 				preservedStructs = cfg.PreservedStructs
 			}
+			if len(preservedFiles) == 0 {
+				preservedFiles = cfg.PreservedFiles
+			}
 		}
 	}
 	forceTrim := false
@@ -108,11 +115,11 @@ func TrimAST(arg *TrimASTArg) (trimResultInfo *TrimResultInfo, err error) {
 	if arg.MatchGoName != nil {
 		matchGoName = *arg.MatchGoName
 	}
-	return doTrimAST(arg.Ast, arg.TrimMethods, forceTrim, matchGoName, preservedStructs)
+	return doTrimAST(arg.Ast, arg.TrimMethods, forceTrim, matchGoName, preservedStructs, preservedFiles)
 }
 
 // doTrimAST trim the single AST, pass method names if -m specified
-func doTrimAST(ast *parser.Thrift, trimMethods []string, forceTrimming, matchGoName bool, preservedStructs []string) (
+func doTrimAST(ast *parser.Thrift, trimMethods []string, forceTrimming, matchGoName bool, preservedStructs, preserveFiles []string) (
 	trimResultInfo *TrimResultInfo, err error) {
 	trimmer, err := newTrimmer(nil, "")
 	if err != nil {
@@ -143,6 +150,7 @@ func doTrimAST(ast *parser.Thrift, trimMethods []string, forceTrimming, matchGoN
 	trimmer.countStructs(ast)
 	originStructsNum := trimmer.structsTrimmed
 	originFieldNum := trimmer.fieldsTrimmed
+	trimmer.loadPreserveFiles(ast, preserveFiles)
 	trimmer.markAST(ast)
 	trimmer.traversal(ast, ast.Filename)
 	if path := parser.CircleDetect(ast); len(path) > 0 {