From 72a7b2eca08add724f115fe52c72e52245583b58 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <stratonikov@runbox.com>
Date: Fri, 17 Mar 2023 18:57:11 +0300
Subject: [PATCH] cmd: allow to provide custom functions for export

As an example for the frostfs-node:
```
go run ./cmd -funcs reportFlushError,reportError ../node
```

Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>
---
 cmd/main.go | 34 +++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index 9a7edfe..201d9ef 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"go/ast"
 	"go/parser"
@@ -8,14 +9,21 @@ import (
 	"log"
 	"os"
 	"path/filepath"
+	"strings"
+)
+
+var (
+	customLogs = flag.String("funcs", "", "Custom functions to export delimited by ','.")
 )
 
 func main() {
-	if len(os.Args) != 2 {
+	flag.Parse()
+
+	if flag.NArg() != 1 {
 		log.Fatalln("directory must be provided")
 	}
 
-	directories, err := getAllSubdirectories(os.Args[1])
+	directories, err := getAllSubdirectories(flag.Arg(0))
 	if err != nil {
 		log.Fatalf("failed to list subdirs: %v", err)
 	}
@@ -78,7 +86,27 @@ func isLogDot(expr ast.Expr) bool {
 	return ok && (isIdent(sel.Sel, "Debug") ||
 		isIdent(sel.Sel, "Info") ||
 		isIdent(sel.Sel, "Warn") ||
-		isIdent(sel.Sel, "Error"))
+		isIdent(sel.Sel, "Error") ||
+		isCustomLog(sel.Sel))
+}
+
+func isCustomLog(expr ast.Expr) bool {
+	id, ok := expr.(*ast.Ident)
+	if !ok {
+		return false
+	}
+
+	if len(*customLogs) == 0 {
+		return false
+	}
+
+	ss := strings.Split(*customLogs, ",")
+	for i := range ss {
+		if id.Name == ss[i] {
+			return true
+		}
+	}
+	return false
 }
 
 func isIdent(expr ast.Expr, ident string) bool {