1//===----------------------------------------------------------------------===//
2//
3//                     The LLVM Compiler Infrastructure
4//
5// This file is dual licensed under the MIT and the University of Illinois Open
6// Source Licenses. See LICENSE.TXT for details.
7//
8//===----------------------------------------------------------------------===//
9
10#ifndef COUNT_NEW_HPP
11#define COUNT_NEW_HPP
12
13# include <cstdlib>
14# include <cassert>
15# include <new>
16
17#ifndef __has_feature
18#  define __has_feature(x) 0
19#endif
20
21#if  __has_feature(address_sanitizer) \
22  || __has_feature(memory_sanitizer) \
23  || __has_feature(thread_sanitizer)
24#define DISABLE_NEW_COUNT
25#endif
26
27class MemCounter
28{
29public:
30    // Make MemCounter super hard to accidentally construct or copy.
31    class MemCounterCtorArg_ {};
32    explicit MemCounter(MemCounterCtorArg_) { reset(); }
33
34private:
35    MemCounter(MemCounter const &);
36    MemCounter & operator=(MemCounter const &);
37
38public:
39    // All checks return true when disable_checking is enabled.
40    static const bool disable_checking;
41
42    // Disallow any allocations from occurring. Useful for testing that
43    // code doesn't perform any allocations.
44    bool disable_allocations;
45
46    int outstanding_new;
47    int new_called;
48    int delete_called;
49    int last_new_size;
50
51    int outstanding_array_new;
52    int new_array_called;
53    int delete_array_called;
54    int last_new_array_size;
55
56public:
57    void newCalled(std::size_t s)
58    {
59        assert(disable_allocations == false);
60        assert(s);
61        ++new_called;
62        ++outstanding_new;
63        last_new_size = s;
64    }
65
66    void deleteCalled(void * p)
67    {
68        assert(p);
69        --outstanding_new;
70        ++delete_called;
71    }
72
73    void newArrayCalled(std::size_t s)
74    {
75        assert(disable_allocations == false);
76        assert(s);
77        ++outstanding_array_new;
78        ++new_array_called;
79        last_new_array_size = s;
80    }
81
82    void deleteArrayCalled(void * p)
83    {
84        assert(p);
85        --outstanding_array_new;
86        ++delete_array_called;
87    }
88
89    void disableAllocations()
90    {
91        disable_allocations = true;
92    }
93
94    void enableAllocations()
95    {
96        disable_allocations = false;
97    }
98
99    void reset()
100    {
101        disable_allocations = false;
102
103        outstanding_new = 0;
104        new_called = 0;
105        delete_called = 0;
106        last_new_size = 0;
107
108        outstanding_array_new = 0;
109        new_array_called = 0;
110        delete_array_called = 0;
111        last_new_array_size = 0;
112    }
113
114public:
115    bool checkOutstandingNewEq(int n) const
116    {
117        return disable_checking || n == outstanding_new;
118    }
119
120    bool checkOutstandingNewNotEq(int n) const
121    {
122        return disable_checking || n != outstanding_new;
123    }
124
125    bool checkNewCalledEq(int n) const
126    {
127        return disable_checking || n == new_called;
128    }
129
130    bool checkNewCalledNotEq(int n) const
131    {
132        return disable_checking || n != new_called;
133    }
134
135    bool checkDeleteCalledEq(int n) const
136    {
137        return disable_checking || n == delete_called;
138    }
139
140    bool checkDeleteCalledNotEq(int n) const
141    {
142        return disable_checking || n != delete_called;
143    }
144
145    bool checkLastNewSizeEq(int n) const
146    {
147        return disable_checking || n == last_new_size;
148    }
149
150    bool checkLastNewSizeNotEq(int n) const
151    {
152        return disable_checking || n != last_new_size;
153    }
154
155    bool checkOutstandingArrayNewEq(int n) const
156    {
157        return disable_checking || n == outstanding_array_new;
158    }
159
160    bool checkOutstandingArrayNewNotEq(int n) const
161    {
162        return disable_checking || n != outstanding_array_new;
163    }
164
165    bool checkNewArrayCalledEq(int n) const
166    {
167        return disable_checking || n == new_array_called;
168    }
169
170    bool checkNewArrayCalledNotEq(int n) const
171    {
172        return disable_checking || n != new_array_called;
173    }
174
175    bool checkDeleteArrayCalledEq(int n) const
176    {
177        return disable_checking || n == delete_array_called;
178    }
179
180    bool checkDeleteArrayCalledNotEq(int n) const
181    {
182        return disable_checking || n != delete_array_called;
183    }
184
185    bool checkLastNewArraySizeEq(int n) const
186    {
187        return disable_checking || n == last_new_array_size;
188    }
189
190    bool checkLastNewArraySizeNotEq(int n) const
191    {
192        return disable_checking || n != last_new_array_size;
193    }
194};
195
196#ifdef DISABLE_NEW_COUNT
197  const bool MemCounter::disable_checking = true;
198#else
199  const bool MemCounter::disable_checking = false;
200#endif
201
202MemCounter globalMemCounter((MemCounter::MemCounterCtorArg_()));
203
204#ifndef DISABLE_NEW_COUNT
205void* operator new(std::size_t s) throw(std::bad_alloc)
206{
207    globalMemCounter.newCalled(s);
208    return std::malloc(s);
209}
210
211void  operator delete(void* p) throw()
212{
213    globalMemCounter.deleteCalled(p);
214    std::free(p);
215}
216
217
218void* operator new[](std::size_t s) throw(std::bad_alloc)
219{
220    globalMemCounter.newArrayCalled(s);
221    return operator new(s);
222}
223
224
225void operator delete[](void* p) throw()
226{
227    globalMemCounter.deleteArrayCalled(p);
228    operator delete(p);
229}
230
231#endif // DISABLE_NEW_COUNT
232
233
234struct DisableAllocationGuard {
235    explicit DisableAllocationGuard(bool disable = true) : m_disabled(disable)
236    {
237        // Don't re-disable if already disabled.
238        if (globalMemCounter.disable_allocations == true) m_disabled = false;
239        if (m_disabled) globalMemCounter.disableAllocations();
240    }
241
242    void release() {
243        if (m_disabled) globalMemCounter.enableAllocations();
244        m_disabled = false;
245    }
246
247    ~DisableAllocationGuard() {
248        release();
249    }
250
251private:
252    bool m_disabled;
253
254    DisableAllocationGuard(DisableAllocationGuard const&);
255    DisableAllocationGuard& operator=(DisableAllocationGuard const&);
256};
257
258#endif /* COUNT_NEW_HPP */
259