-
Notifications
You must be signed in to change notification settings - Fork 151
/
pool.go
130 lines (115 loc) · 3.36 KB
/
pool.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 2020 ConsenSys Software 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.
// Code generated by consensys/gnark-crypto DO NOT EDIT
package polynomial
import (
"fmt"
"github.com/consensys/gnark-crypto/ecc/bw6-633/fr"
"reflect"
"sync"
"unsafe"
)
// Memory management for polynomials
// Copied verbatim from gkr repo
// Sets a maximum for the array size we keep in pool
const maxNForLargePool int = 1 << 24
const maxNForSmallPool int = 256
// Aliases because it is annoying to use arrays in all the places
type largeArr = [maxNForLargePool]fr.Element
type smallArr = [maxNForSmallPool]fr.Element
var rC = sync.Map{}
var (
largePool = sync.Pool{
New: func() interface{} {
var res largeArr
return &res
},
}
smallPool = sync.Pool{
New: func() interface{} {
var res smallArr
return &res
},
}
)
// ClearPool Clears the pool completely, shields against memory leaks
// Eg: if we forgot to dump a polynomial at some point, this will ensure the value get dumped eventually
// Returns how many polynomials were cleared that way
func ClearPool() int {
res := 0
rC.Range(func(k, _ interface{}) bool {
switch ptr := k.(type) {
case *largeArr:
largePool.Put(ptr)
case *smallArr:
smallPool.Put(ptr)
default:
panic(fmt.Sprintf("tried to clear %v", reflect.TypeOf(ptr)))
}
res++
return true
})
return res
}
// CountPool Returns the number of elements in the pool without mutating it
func CountPool() int {
res := 0
rC.Range(func(_, _ interface{}) bool {
res++
return true
})
return res
}
// Make tries to find a reusable polynomial or allocates a new one
func Make(n int) []fr.Element {
if n > maxNForLargePool {
panic(fmt.Sprintf("been provided with size of %v but the maximum is %v", n, maxNForLargePool))
}
if n <= maxNForSmallPool {
ptr := smallPool.Get().(*smallArr)
rC.Store(ptr, struct{}{}) // registers the pointer being used
return (*ptr)[:n]
}
ptr := largePool.Get().(*largeArr)
rC.Store(ptr, struct{}{}) // remember we allocated the pointer is being used
return (*ptr)[:n]
}
// Dump dumps a set of polynomials into the pool
// Returns the number of deallocated polys
func Dump(arrs ...[]fr.Element) int {
cnt := 0
for _, arr := range arrs {
ptr := ptr(arr)
pool := &smallPool
if len(arr) > maxNForSmallPool {
pool = &largePool
}
// If the rC did not register, then
// either the array was allocated somewhere else which can be ignored
// otherwise a double put which MUST be ignored
if _, ok := rC.Load(ptr); ok {
pool.Put(ptr)
// And deregisters the ptr
rC.Delete(ptr)
cnt++
}
}
return cnt
}
func ptr(m []fr.Element) unsafe.Pointer {
if cap(m) != maxNForSmallPool && cap(m) != maxNForLargePool {
panic(fmt.Sprintf("can't cast to large or small array, the put array's is %v it should have capacity %v or %v", cap(m), maxNForLargePool, maxNForSmallPool))
}
return unsafe.Pointer(&m[0])
}