forked from google/agi
-
Notifications
You must be signed in to change notification settings - Fork 0
/
types.go
130 lines (118 loc) · 3.79 KB
/
types.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
// Copyright (C) 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pack
import (
"context"
"fmt"
"reflect"
"github.com/golang/protobuf/proto"
"github.com/golang/protobuf/protoc-gen-go/descriptor"
"github.com/google/gapid/core/data/protoutil"
)
// ty is an entry in the map of types stored in a packfile.
type ty struct {
// name is the cannocial unique name of the type.
// This will be the same name as used in the proto registry.
name string
// index is the tag index used for the type in this packfile.
index uint64
// create constructs a new proto of this type.
create func() proto.Message
// desc is the proto description of this type, it is packed
// into the file and can be used to reflect on the type.
desc *descriptor.DescriptorProto
}
// types stores the full type registry for a packfile.
// It is exposed so that you can pre-build a cannocial type registry
// rather than constructing on demand.
type types struct {
entries []*ty
byName map[string]*ty
forceDynamic bool
}
// newTypes constructs a new empty type registry.
func newTypes(forceDynamic bool) *types {
return &types{
entries: []*ty{nil},
byName: map[string]*ty{},
forceDynamic: forceDynamic,
}
}
// addForMessage adds all types required to serialize the message.
// The callback will be called for each newly added type.
func (t *types) addForMessage(ctx context.Context, msg proto.Message, cb func(t *ty) error) (*ty, error) {
return t.addForType(ctx, reflect.TypeOf(msg), cb)
}
// addForType adds all types required to serialize the reflection type.
// The callback will be called for each newly added type.
func (t *types) addForType(ctx context.Context, typ reflect.Type, cb func(t *ty) error) (*ty, error) {
for typ.Kind() == reflect.Ptr {
typ = typ.Elem()
}
msg := reflect.New(typ).Interface().(proto.Message)
name := proto.MessageName(msg)
if name == "" {
panic(fmt.Errorf("Can not determine the message name of: %v", msg))
}
if ty, ok := t.byName[name]; ok {
return ty, nil
}
var desc *descriptor.DescriptorProto
if d, ok := msg.(protoutil.Described); ok {
desc, _ = protoutil.DescriptorOf(d)
}
ty := t.add(name, desc)
if err := cb(ty); err != nil {
return nil, err
}
// Recursively add referenced types.
if typ.Kind() == reflect.Struct {
numFields := typ.NumField()
protoMsgTy := reflect.TypeOf((*proto.Message)(nil)).Elem()
for i := 0; i < numFields; i++ {
fieldType := typ.Field(i).Type
for fieldType.Kind() == reflect.Slice {
fieldType = fieldType.Elem()
}
if fieldType.Implements(protoMsgTy) {
if _, err := t.addForType(ctx, fieldType, cb); err != nil {
return nil, err
}
}
}
}
return ty, nil
}
// count returns the number of types in the registry.
func (t *types) count() uint64 {
return uint64(len(t.entries))
}
// add adds a type by name and descriptor.
func (t *types) add(name string, desc *descriptor.DescriptorProto) *ty {
create := func() proto.Message { return newDynamic(desc, t) }
if !t.forceDynamic {
if ty := proto.MessageType(name); ty != nil {
create = func() proto.Message { return reflect.New(ty.Elem()).Interface().(proto.Message) }
}
}
entry := &ty{
name: name,
index: t.count(),
create: create,
desc: desc,
}
t.entries = append(t.entries, entry)
t.byName[name] = entry
return entry
}