Code generator: big map subscriptions

.vscode/launch.json

@@ -11,7 +11,7 @@
             "mode": "debug",
             "program": "${workspaceFolder}/cmd/generator",
             "args": [
-                "-n", "atomex", "-c", "KT1VG2WtYdSWz5E7chTeAdDPZNy2MpP8pTfL"
+                "-n", "atomex", "-c", "KT1KHdx4EC7WNVeeTGMr8utCSRczgu2EPCzX", "-u", "https://api.granadanet.tzkt.io"
             ]
         }
     ]

cmd/generator/README.md

@@ -84,6 +84,8 @@ func main() {
             log.Println(redeem)
         case refund := <-atx.RefundEvents():
             log.Println(refund)
+        case update := <-atx.BigMap0Updates():
+            log.Println(update)
         }
     }
 }

cmd/generator/generator.go

@@ -17,6 +17,7 @@ type templateContext struct {
 	PackageName     string
 	TypeName        string
 	EntrypointTypes map[string]types.EntrypointData
+	BigMaps         map[string]types.BigMapData
 	Contract        string
 }
 
@@ -64,6 +65,7 @@ func generateContractObject(name, contract, dest string, result types.ContractTy
 		TypeName:        className,
 		Contract:        contract,
 		EntrypointTypes: result.Entrypoints,
+		BigMaps:         result.BigMaps,
 	})
 }
 
@@ -89,7 +91,12 @@ func generateDefaultTypes(packageName, dest string) error {
 }
 
 func generateFromTemplate(templateFileName, dest string, ctx interface{}) error {
-	tmpl, err := template.ParseFS(templates, "template/*")
+	tmpl, err := template.New("").Funcs(template.FuncMap{
+		"ToLower": strings.ToLower,
+		"TrimStorage": func(x string) string {
+			return strings.TrimPrefix(x, "storage.")
+		},
+	}).ParseFS(templates, "template/*")
 	if err != nil {
 		return err
 	}

cmd/generator/template/contract.tmpl

@@ -24,10 +24,12 @@ type {{ $value.Type }}Tx struct {
 type {{.TypeName}} struct {
 	tzktAPI *api.API
 	tzktEvents *events.TzKT
-
-{{- range $key, $value := .EntrypointTypes }}
+{{ range $key, $value := .EntrypointTypes }}
 	{{ $value.Var }} chan {{ $value.Type }}Tx
 {{- end }}
+{{- range $key, $value := .BigMaps }}
+	{{ $value.Type | ToLower }} chan {{ $value.Type }}
+{{ end }}
 	wg sync.WaitGroup
 }
 
@@ -39,6 +41,9 @@ func New(baseURL string) *{{.TypeName}}  {
 		{{- range $key, $value := .EntrypointTypes }}
 		{{ $value.Var }}: make(chan {{ $value.Type }}Tx, 1024),
 		{{- end }}
+		{{- range $key, $value := .BigMaps }}
+		{{ $value.Type | ToLower }}: make(chan {{ $value.Type }}, 1024),
+		{{- end }}
 	}
 }
 
@@ -51,6 +56,10 @@ func (contract *{{.TypeName}}) Subscribe(ctx context.Context) error {
 	contract.wg.Add(1)
 	go contract.listen(ctx)
 
+	if err := contract.tzktEvents.SubscribeToBigMaps(nil, "{{ .Contract }}", ""); err != nil {
+		return err
+	}
+
 	return contract.tzktEvents.SubscribeToOperations("{{ .Contract }}", api.KindTransaction)
 }
 
@@ -74,6 +83,12 @@ func (contract *{{$.TypeName}}) {{$value.Type}}Events() <-chan {{ $value.Type }}
 	return contract.{{$value.Var}}
 }
 {{ end }}
+{{ range $key, $value := .BigMaps }}
+// {{$value.Type}}Updates - listen `{{$value.Type}}` updates channel
+func (contract *{{$.TypeName}}) {{$value.Type}}Updates() <-chan {{ $value.Type }} {
+	return contract.{{$value.Type | ToLower }}
+}
+{{ end }}
 
 func (contract *{{.TypeName}}) listen(ctx context.Context) {
 	defer contract.wg.Done()
@@ -87,6 +102,31 @@ func (contract *{{.TypeName}}) listen(ctx context.Context) {
 			case events.MessageTypeData:
 
 				switch msg.Channel {
+				case events.ChannelBigMap:
+					items := msg.Body.([]events.BigMapUpdate)
+					for i := range items {
+						switch items[i].Path {
+						{{- range $key, $value := .BigMaps }}
+						case "{{$key | TrimStorage}}":
+							var key {{$value.KeyType}}
+							if err := json.Unmarshal([]byte(items[i].Content.Key), &key); err != nil {
+								log.Println(err)
+								continue
+							}
+
+							var value {{$value.ValueType}}
+							if err := json.Unmarshal(items[i].Content.Value, &value); err != nil {
+								log.Println(err)
+								continue
+							}
+							contract.{{$value.Type | ToLower }} <- {{$value.Type}}{
+								Key: key,
+								Value: value,
+							}
+						{{- end }}
+						}
+					}
+
 				case events.ChannelOperations:
 					items := msg.Body.([]interface{})
 					for _, item := range items {

cmd/generator/types/address.go

@@ -18,12 +18,12 @@ func (Address) AsField(name, path string, schema api.JSONSchema, isRequired bool
 		tags["validate"] = TagRequired
 	}
 
-	return jen.Id(fieldName(name)).Add(jen.Id("Address")).Tag(tags), nil
+	return jen.Id(fieldName("Address", name)).Add(jen.Id("Address")).Tag(tags), nil
 }
 
 // AsCode -
 func (Address) AsCode(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
-	typName := result.GetName(name)
+	typName := result.GetName("Address", name)
 	return Code{
 		Statement: jen.Comment(typName).Line().Type().Id(typName).Add(jen.Id("Address")).Line(),
 		Name:      typName,

cmd/generator/types/big_map.go

@@ -19,7 +19,7 @@ func GenerateBigMap(bigMap api.BigMapJSONSchema, result *ContractTypeResult) err
 		return err
 	}
 
-	typeName := result.GetName(bigMap.Name)
+	typeName := result.GetName("BigMap", bigMap.Name)
 	result.File.Comment(typeName).Line().Type().Id(typeName).Struct(
 		jen.Id("Key").Add(jen.Id(keyTypeName)),
 		jen.Id("Value").Add(jen.Id(valueTypeName)),
@@ -52,7 +52,11 @@ func GenerateBigMap(bigMap api.BigMapJSONSchema, result *ContractTypeResult) err
 		),
 	)
 
-	result.bigMaps[getPath("storage", bigMap.Path)] = typeName
+	result.BigMaps[getPath("storage", bigMap.Path)] = BigMapData{
+		KeyType:   keyTypeName,
+		ValueType: valueTypeName,
+		Type:      typeName,
+	}
 
 	return nil
 }
@@ -70,25 +74,28 @@ func (BigMap) AsField(name, path string, schema api.JSONSchema, isRequired bool,
 		tags["validate"] = TagRequired
 	}
 
-	typName, ok := result.bigMaps[getPath(path, name)]
+	var typName string
+	bmData, ok := result.BigMaps[getPath(path, name)]
 	if !ok {
-		typName = result.GetName(name)
+		typName = result.GetName("BigMap", name)
+	} else {
+		typName = bmData.Type
 	}
 
-	return jen.Id(fieldName(name)).Add(jen.Id(typName)).Tag(tags), nil
+	return jen.Id(fieldName("BigMap", name)).Add(jen.Id(typName)).Tag(tags), nil
 }
 
 // AsCode -
 func (BigMap) AsCode(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
 	return Code{
 		Statement: nil,
-		Name:      result.GetName(name),
+		Name:      result.GetName("BigMap", name),
 	}, nil
 }
 
 // AsType -
 func (BigMap) AsType(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
-	typName := result.GetName(name)
+	typName := result.GetName("BigMap", name)
 	return Code{
 		Statement: jen.Id(typName),
 		Name:      typName,

cmd/generator/types/bool.go

@@ -18,12 +18,12 @@ func (Bool) AsField(name, path string, schema api.JSONSchema, isRequired bool, r
 		tags["validate"] = TagRequired
 	}
 
-	return jen.Id(fieldName(name)).Bool().Tag(tags), nil
+	return jen.Id(fieldName("Bool", name)).Bool().Tag(tags), nil
 }
 
 // AsCode -
 func (Bool) AsCode(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
-	typ := result.GetName(name)
+	typ := result.GetName("Bool", name)
 	return Code{
 		Statement: jen.Comment(typ).Line().Type().Id(typ).Bool().Line(),
 		Name:      typ,
@@ -34,6 +34,6 @@ func (Bool) AsCode(name, path string, schema api.JSONSchema, result *ContractTyp
 func (Bool) AsType(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
 	return Code{
 		Statement: jen.Bool(),
-		Name:      result.GetName(name),
+		Name:      result.GetName("Bool", name),
 	}, nil
 }

cmd/generator/types/bytes.go

@@ -18,12 +18,12 @@ func (Bytes) AsField(name, path string, schema api.JSONSchema, isRequired bool,
 		tags["validate"] = TagRequired
 	}
 
-	return jen.Id(fieldName(name)).Add(jen.Id("Bytes")).Tag(tags), nil
+	return jen.Id(fieldName("Bytes", name)).Add(jen.Id("Bytes")).Tag(tags), nil
 }
 
 // AsCode -
 func (Bytes) AsCode(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
-	typName := result.GetName(name)
+	typName := result.GetName("Bytes", name)
 	return Code{
 		Statement: jen.Comment(typName).Line().Type().Id(typName).Add(jen.Id("Bytes")).Line(),
 		Name:      typName,
@@ -34,6 +34,6 @@ func (Bytes) AsCode(name, path string, schema api.JSONSchema, result *ContractTy
 func (Bytes) AsType(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
 	return Code{
 		Statement: jen.Add(jen.Id("Bytes")),
-		Name:      result.GetName(name),
+		Name:      result.GetName("Bytes", name),
 	}, nil
 }

cmd/generator/types/contract.go

@@ -18,12 +18,12 @@ func (Contract) AsField(name, path string, schema api.JSONSchema, isRequired boo
 		tags["validate"] = TagRequired
 	}
 
-	return jen.Id(fieldName(name)).Add(jen.Id("Contract")).Tag(tags), nil
+	return jen.Id(fieldName("Contract", name)).Add(jen.Id("Contract")).Tag(tags), nil
 }
 
 // AsCode -
 func (Contract) AsCode(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
-	typName := result.GetName(name)
+	typName := result.GetName("Contract", name)
 	return Code{
 		Statement: jen.Comment(typName).Line().Type().Id(typName).Add(jen.Id("Contract")).Line(),
 		Name:      typName,

cmd/generator/types/general.go

@@ -3,6 +3,7 @@ package types
 import (
 	"fmt"
 	"go/token"
+	"strconv"
 
 	"github.com/dave/jennifer/jen"
 	"github.com/dipdup-net/go-lib/tzkt/api"
@@ -51,24 +52,24 @@ func Generate(name string, schema api.JSONSchema, result *ContractTypeResult) (s
 type ContractTypeResult struct {
 	File        *jen.File
 	Entrypoints map[string]EntrypointData
+	BigMaps     map[string]BigMapData
 	PackageName string
 
 	names   map[string]struct{}
-	bigMaps map[string]string
 	counter int64
 }
 
 // GetName -
-func (result *ContractTypeResult) GetName(name string) string {
+func (result *ContractTypeResult) GetName(typ, name string) string {
 	if reserved, ok := reservedNames[name]; ok {
 		name = reserved
 	}
-	name = strcase.ToCamel(name)
+	name = fieldName(typ, name)
 
 	if _, exists := result.names[name]; exists {
 		result.counter++
 		name = fmt.Sprintf("%s%d", name, result.counter)
-		name = result.GetName(name)
+		name = result.GetName(typ, name)
 	}
 
 	result.names[name] = struct{}{}
@@ -81,15 +82,22 @@ type EntrypointData struct {
 	Var  string
 }
 
+// BigMapData -
+type BigMapData struct {
+	Type      string
+	KeyType   string
+	ValueType string
+}
+
 // GenerateContractTypes -
 func GenerateContractTypes(schema api.ContractJSONSchema, packageName string) (ContractTypeResult, error) {
 	result := ContractTypeResult{
 		File:        jen.NewFile(packageName),
 		Entrypoints: make(map[string]EntrypointData),
+		BigMaps:     make(map[string]BigMapData),
 		PackageName: packageName,
 
-		names:   make(map[string]struct{}),
-		bigMaps: make(map[string]string),
+		names: make(map[string]struct{}),
 	}
 
 	result.File.PackageComment("DO NOT EDIT!!!")
@@ -206,7 +214,10 @@ func isSimpleType(comment string) bool {
 	return false
 }
 
-func fieldName(name string) string {
+func fieldName(typ, name string) string {
+	if _, err := strconv.ParseInt(name, 10, 64); err == nil {
+		name = fmt.Sprintf("%s%s", typ, name)
+	}
 	return strcase.ToCamel(name)
 }
 

cmd/generator/types/int.go

@@ -20,12 +20,12 @@ func (Int) AsField(name, path string, schema api.JSONSchema, isRequired bool, re
 		tags["validate"] = TagRequired
 	}
 
-	return jen.Id(fieldName(name)).Int64().Tag(tags), nil
+	return jen.Id(fieldName("Int", name)).Int64().Tag(tags), nil
 }
 
 // AsCode -
 func (Int) AsCode(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
-	typ := result.GetName(name)
+	typ := result.GetName("Int", name)
 	return Code{
 		Statement: jen.Comment(typ).Line().Type().Id(typ).Int64().Line(),
 		Name:      typ,
@@ -36,6 +36,6 @@ func (Int) AsCode(name, path string, schema api.JSONSchema, result *ContractType
 func (Int) AsType(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
 	return Code{
 		Statement: jen.Int64(),
-		Name:      result.GetName(name),
+		Name:      result.GetName("Int", name),
 	}, nil
 }

cmd/generator/types/key_hash.go

@@ -18,12 +18,12 @@ func (KeyHash) AsField(name, path string, schema api.JSONSchema, isRequired bool
 		tags["validate"] = TagRequired
 	}
 
-	return jen.Id(fieldName(name)).String().Tag(tags), nil
+	return jen.Id(fieldName("KeyHash", name)).String().Tag(tags), nil
 }
 
 // AsCode -
 func (KeyHash) AsCode(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
-	typName := result.GetName(name)
+	typName := result.GetName("KeyHash", name)
 	return Code{
 		Statement: jen.Comment(typName).Line().Type().Id(typName).String().Line(),
 		Name:      typName,

cmd/generator/types/lambda.go

@@ -18,12 +18,12 @@ func (Lambda) AsField(name, path string, schema api.JSONSchema, isRequired bool,
 		tags["validate"] = TagRequired
 	}
 
-	return jen.Id(fieldName(name)).String().Tag(tags), nil
+	return jen.Id(fieldName("Lambda", name)).String().Tag(tags), nil
 }
 
 // AsCode -
 func (Lambda) AsCode(name, path string, schema api.JSONSchema, result *ContractTypeResult) (Code, error) {
-	typName := result.GetName(name)
+	typName := result.GetName("Lambda", name)
 	return Code{
 		Statement: jen.Comment(typName).Line().Type().Id(typName).String().Line(),
 		Name:      typName,