From ede37de7133084fa1a01074ded0d74c9f0a75d25 Mon Sep 17 00:00:00 2001
From: HeyJavaBean <liyun.339@bytedance.com>
Date: Fri, 13 Dec 2024 14:30:27 +0800
Subject: [PATCH 1/2] optimize: trimmer config idl file preserve

---
 tool/trimmer/trim/config.go  |  1 +
 tool/trimmer/trim/mark.go    | 24 +++++++++++++++++++++++-
 tool/trimmer/trim/trimmer.go | 13 ++++++++++---
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/tool/trimmer/trim/config.go b/tool/trimmer/trim/config.go
index a88a497..40b4ab8 100644
--- a/tool/trimmer/trim/config.go
+++ b/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 {
diff --git a/tool/trimmer/trim/mark.go b/tool/trimmer/trim/mark.go
index b4f0c21..9f2bae9 100644
--- a/tool/trimmer/trim/mark.go
+++ b/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
+			}
+		}
+	}
 }
diff --git a/tool/trimmer/trim/trimmer.go b/tool/trimmer/trim/trimmer.go
index bdf22cb..be34689 100644
--- a/tool/trimmer/trim/trimmer.go
+++ b/tool/trimmer/trim/trimmer.go
@@ -45,6 +45,7 @@ type Trimmer struct {
 	structsTrimmed   int
 	fieldsTrimmed    int
 	extServices      []*parser.Service
+	PreservedFiles   []string
 }
 
 type TrimASTArg struct {
@@ -53,6 +54,7 @@ type TrimASTArg struct {
 	Preserve        *bool
 	MatchGoName     *bool
 	PreserveStructs []string
+	PreservedFiles  []string
 }
 
 type TrimResultInfo struct {
@@ -80,8 +82,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 +101,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 +114,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 +149,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 {

From 7b233e0cde25b378cb5b6857a2d5587eb3a45c3f Mon Sep 17 00:00:00 2001
From: HeyJavaBean <liyun.339@bytedance.com>
Date: Fri, 13 Dec 2024 14:32:52 +0800
Subject: [PATCH 2/2] fix: fix bug

---
 tool/trimmer/trim/trimmer.go | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/tool/trimmer/trim/trimmer.go b/tool/trimmer/trim/trimmer.go
index be34689..a4c2864 100644
--- a/tool/trimmer/trim/trimmer.go
+++ b/tool/trimmer/trim/trimmer.go
@@ -36,16 +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
-	PreservedFiles   []string
+	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 {