parse for all the structs

1 files changed, 28 insertions, 8 deletions

diff --git a/parseProtoFile.go b/parseProtoFile.go
index cbe670a..2795bbd 100644
--- a/parseProtoFile.go
+++ b/parseProtoFile.go
@@ -25,6 +25,13 @@ func (pb *Files) findAutogenpb(f *File) error {
 		return err
 	}
 
+	// first parse the proto file for message struct names
+	for _, line := range strings.Split(string(data), "\n") {
+		if strings.HasPrefix(line, "message ") {
+			f.parseForMessage(line)
+		}
+	}
+
 	// look for included proto files
 	lines := strings.Split(string(data), "\n")
 	for _, line := range lines {
@@ -41,18 +48,31 @@ func (pb *Files) findAutogenpb(f *File) error {
 			// log.Info("found unique field", newu)
 			uniqueKeys = append(uniqueKeys, newu)
 		}
-		if strings.Contains(line, "autogenpb:mutex") {
-			parts := strings.Split(line, "autogenpb:mutex")
-			// log.Info("FOUND MUTEX line:", parts[0])
-			fields := strings.Fields(parts[0])
-			if fields[0] == "message" {
-				log.Info("FOUND MUTEX:", fields[1])
-			}
-		}
 	}
 	return nil
 }
 
+// looks for mutex and marshal entries
+func (f *File) parseForMessage(line string) {
+	fields := strings.Fields(line)
+	if fields[0] != "message" {
+		return
+	}
+	msgName := fields[1]
+	msg := new(MsgName)
+	f.MsgNames = append(f.MsgNames, msg)
+	msg.Name = msgName
+
+	if strings.Contains(line, "`autogenpb:mutex`") {
+		msg.DoMutex = true
+		log.Info("Found Mutex for:", msg.Name)
+	}
+	if strings.Contains(line, "`autogenpb:marshal`") {
+		msg.DoMarshal = true
+		log.Info("Found Marshal for:", msg.Name)
+	}
+}
+
 func (pb *Files) findGlobalAutogenpb(f *File) error {
 	// log.Info("starting findAutogenpb() on", filename)
 	// read in the .proto file