From 5267e8aaecd0cdcd0e15f8f9cdb59c08dc473e5b Mon Sep 17 00:00:00 2001
From: Tolya Korniltsev <korniltsev.anatoly@gmail.com>
Date: Fri, 10 Jun 2022 22:03:15 +0700
Subject: [PATCH] cpool parsing processor option

---
 parser/checkpoint.go |  4 ++--
 parser/chunk.go      | 16 +++++++++++++---
 parser/parser.go     |  6 +++++-
 parser/types.go      |  4 ++--
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/parser/checkpoint.go b/parser/checkpoint.go
index dfef7d0..7a4c8c4 100644
--- a/parser/checkpoint.go
+++ b/parser/checkpoint.go
@@ -41,7 +41,7 @@ func (c *CheckpointEvent) Parse(r reader.Reader, classes ClassMap, cpools PoolMa
 		}
 		cm, ok := cpools[int(classID)]
 		if !ok {
-			cpools[int(classID)] = &CPool{pool: make(map[int]ParseResolvable)}
+			cpools[int(classID)] = &CPool{Pool: make(map[int]ParseResolvable)}
 			cm = cpools[int(classID)]
 		}
 		m, err := r.VarInt()
@@ -58,7 +58,7 @@ func (c *CheckpointEvent) Parse(r reader.Reader, classes ClassMap, cpools PoolMa
 			if err != nil {
 				return fmt.Errorf("unable to parse constant type %d: %w", classID, err)
 			}
-			cm.pool[int(idx)] = v
+			cm.Pool[int(idx)] = v
 		}
 	}
 	return nil
diff --git a/parser/chunk.go b/parser/chunk.go
index 320b21e..fbf2e72 100644
--- a/parser/chunk.go
+++ b/parser/chunk.go
@@ -11,7 +11,7 @@ import (
 var magic = []byte{'F', 'L', 'R', 0}
 
 type CPool struct {
-	pool     map[int]ParseResolvable
+	Pool     map[int]ParseResolvable
 	resolved bool
 }
 type ClassMap map[int]ClassMetadata
@@ -24,7 +24,11 @@ type Chunk struct {
 	Events      []Parseable
 }
 
-func (c *Chunk) Parse(r io.Reader) (err error) {
+type ChunkParseOptions struct {
+	CPoolProcessor func(meta ClassMetadata, cpool *CPool)
+}
+
+func (c *Chunk) Parse(r io.Reader, options *ChunkParseOptions) (err error) {
 	buf := make([]byte, len(magic))
 	if _, err = io.ReadFull(r, buf); err != nil {
 		if err == io.EOF {
@@ -103,6 +107,12 @@ func (c *Chunk) Parse(r io.Reader) (err error) {
 		br.Seek(c.Header.ConstantPoolOffset+delta, io.SeekStart)
 	}
 
+	if options.CPoolProcessor != nil {
+		for classID, pool := range cpools {
+			options.CPoolProcessor(classes[classID], pool)
+		}
+	}
+
 	// Second pass over constant pools: resolve constants
 	for classID := range cpools {
 		if err := ResolveConstants(classes, cpools, classID); err != nil {
@@ -153,7 +163,7 @@ func ResolveConstants(classes ClassMap, cpools PoolMap, classID int) (err error)
 		return nil
 	}
 	cpool.resolved = true
-	for _, t := range cpool.pool {
+	for _, t := range cpool.Pool {
 		if err := t.Resolve(classes, cpools); err != nil {
 			return fmt.Errorf("unable to resolve constants: %w", err)
 		}
diff --git a/parser/parser.go b/parser/parser.go
index 651f4da..069955c 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -6,10 +6,14 @@ import (
 )
 
 func Parse(r io.Reader) ([]Chunk, error) {
+	return ParseWithOptions(r, &ChunkParseOptions{})
+}
+
+func ParseWithOptions(r io.Reader, options *ChunkParseOptions) ([]Chunk, error) {
 	var chunks []Chunk
 	for {
 		var chunk Chunk
-		err := chunk.Parse(r)
+		err := chunk.Parse(r, options)
 		if err == io.EOF {
 			return chunks, nil
 		}
diff --git a/parser/types.go b/parser/types.go
index 2d36b3b..7db6f70 100644
--- a/parser/types.go
+++ b/parser/types.go
@@ -97,7 +97,7 @@ func parseFields(r reader.Reader, classes ClassMap, cpools PoolMap, class ClassM
 				if err != nil {
 					return fmt.Errorf("unable to read constant index")
 				}
-				p, ok := cpool.pool[int(i)]
+				p, ok := cpool.Pool[int(i)]
 				if !ok {
 					continue
 				}
@@ -147,7 +147,7 @@ func resolveConstants(classes ClassMap, cpools PoolMap, constants *[]constant, r
 			// Non-existent constant pool references seem to be used to mark no value
 			continue
 		}
-		it, ok := p.pool[int(c.index)]
+		it, ok := p.Pool[int(c.index)]
 		if !ok {
 			// Non-existent constant pool references seem to be used to mark no value
 			continue