47aedd428f79e5358e0feb83978a661cd55d5451
[wimlib] / src / util.c
1 /*
2  * util.c - utility functions
3  */
4
5 /*
6  * Copyright (C) 2012-2016 Eric Biggers
7  *
8  * This file is free software; you can redistribute it and/or modify it under
9  * the terms of the GNU Lesser General Public License as published by the Free
10  * Software Foundation; either version 3 of the License, or (at your option) any
11  * later version.
12  *
13  * This file is distributed in the hope that it will be useful, but WITHOUT
14  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
15  * FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
16  * details.
17  *
18  * You should have received a copy of the GNU Lesser General Public License
19  * along with this file; if not, see http://www.gnu.org/licenses/.
20  */
21
22 #ifdef HAVE_CONFIG_H
23 #  include "config.h"
24 #endif
25
26 #include <errno.h>
27 #include <limits.h>
28 #include <stdio.h>
29 #include <stdlib.h>
30 #include <string.h>
31 #ifdef HAVE_SYS_SYSCTL_H
32 #  include <sys/types.h>
33 #  include <sys/sysctl.h>
34 #endif
35 #include <unistd.h>
36
37 #include "wimlib.h"
38 #include "wimlib/assert.h"
39 #include "wimlib/error.h"
40 #include "wimlib/timestamp.h"
41 #include "wimlib/util.h"
42 #include "wimlib/xml.h"
43
44 /*******************
45  * Memory allocation
46  *******************/
47
48 static void *(*wimlib_malloc_func) (size_t)         = malloc;
49 static void  (*wimlib_free_func)   (void *)         = free;
50 static void *(*wimlib_realloc_func)(void *, size_t) = realloc;
51
52 void *
53 wimlib_malloc(size_t size)
54 {
55         void *ptr;
56
57 retry:
58         ptr = (*wimlib_malloc_func)(size);
59         if (unlikely(!ptr)) {
60                 if (size == 0) {
61                         size = 1;
62                         goto retry;
63                 }
64         }
65         return ptr;
66 }
67
68 void
69 wimlib_free_memory(void *ptr)
70 {
71         (*wimlib_free_func)(ptr);
72 }
73
74 void *
75 wimlib_realloc(void *ptr, size_t size)
76 {
77         if (size == 0)
78                 size = 1;
79         return (*wimlib_realloc_func)(ptr, size);
80 }
81
82 void *
83 wimlib_calloc(size_t nmemb, size_t size)
84 {
85         size_t total_size = nmemb * size;
86         void *p;
87
88         if (size != 0 && nmemb > SIZE_MAX / size) {
89                 errno = ENOMEM;
90                 return NULL;
91         }
92
93         p = MALLOC(total_size);
94         if (p)
95                 p = memset(p, 0, total_size);
96         return p;
97 }
98
99 char *
100 wimlib_strdup(const char *str)
101 {
102         return memdup(str, strlen(str) + 1);
103 }
104
105 #ifdef __WIN32__
106 wchar_t *
107 wimlib_wcsdup(const wchar_t *str)
108 {
109         return memdup(str, (wcslen(str) + 1) * sizeof(wchar_t));
110 }
111 #endif
112
113 void *
114 wimlib_aligned_malloc(size_t size, size_t alignment)
115 {
116         wimlib_assert(is_power_of_2(alignment));
117
118         void *ptr = MALLOC(sizeof(void *) + alignment - 1 + size);
119         if (ptr) {
120                 void *orig_ptr = ptr;
121                 ptr = (void *)ALIGN((uintptr_t)ptr + sizeof(void *), alignment);
122                 ((void **)ptr)[-1] = orig_ptr;
123         }
124         return ptr;
125 }
126
127 void
128 wimlib_aligned_free(void *ptr)
129 {
130         if (ptr)
131                 FREE(((void **)ptr)[-1]);
132 }
133
134 void *
135 memdup(const void *mem, size_t size)
136 {
137         void *ptr = MALLOC(size);
138         if (ptr)
139                 ptr = memcpy(ptr, mem, size);
140         return ptr;
141 }
142
143 /* API function documented in wimlib.h  */
144 WIMLIBAPI int
145 wimlib_set_memory_allocator(void *(*malloc_func)(size_t),
146                             void (*free_func)(void *),
147                             void *(*realloc_func)(void *, size_t))
148 {
149         wimlib_malloc_func  = malloc_func  ? malloc_func  : malloc;
150         wimlib_free_func    = free_func    ? free_func    : free;
151         wimlib_realloc_func = realloc_func ? realloc_func : realloc;
152
153         xml_set_memory_allocator(wimlib_malloc_func, wimlib_free_func,
154                                  wimlib_realloc_func);
155         return 0;
156 }
157
158 /*******************
159  * String utilities
160  *******************/
161
162 #ifndef HAVE_MEMPCPY
163 void *mempcpy(void *dst, const void *src, size_t n)
164 {
165         return memcpy(dst, src, n) + n;
166 }
167 #endif
168
169 static bool seeded = false;
170
171 static void
172 seed_random(void)
173 {
174         srand(now_as_wim_timestamp());
175         seeded = true;
176 }
177
178 /* Fills @n characters pointed to by @p with random alphanumeric characters. */
179 void
180 randomize_char_array_with_alnum(tchar *p, size_t n)
181 {
182         if (!seeded)
183                 seed_random();
184         while (n--) {
185                 int r = rand() % 62;
186                 if (r < 26)
187                         *p++ = r + 'a';
188                 else if (r < 52)
189                         *p++ = r - 26 + 'A';
190                 else
191                         *p++ = r - 52 + '0';
192         }
193 }
194
195 /* Fills @n bytes pointer to by @p with random numbers. */
196 void
197 randomize_byte_array(u8 *p, size_t n)
198 {
199         if (!seeded)
200                 seed_random();
201         while (n--)
202                 *p++ = rand();
203 }
204
205 #ifndef __WIN32__
206 unsigned
207 get_available_cpus(void)
208 {
209         long n = sysconf(_SC_NPROCESSORS_ONLN);
210         if (n < 1 || n >= UINT_MAX) {
211                 WARNING("Failed to determine number of processors; assuming 1.");
212                 return 1;
213         }
214         return n;
215 }
216 #endif /* !__WIN32__ */
217
218 #ifndef __WIN32__
219 u64
220 get_available_memory(void)
221 {
222 #if defined(_SC_PAGESIZE) && defined(_SC_PHYS_PAGES)
223         long page_size = sysconf(_SC_PAGESIZE);
224         long num_pages = sysconf(_SC_PHYS_PAGES);
225         if (page_size <= 0 || num_pages <= 0)
226                 goto default_size;
227         return ((u64)page_size * (u64)num_pages);
228 #else
229         int mib[2] = {CTL_HW, HW_MEMSIZE};
230         u64 memsize;
231         size_t len = sizeof(memsize);
232         if (sysctl(mib, ARRAY_LEN(mib), &memsize, &len, NULL, 0) < 0 || len != 8)
233                 goto default_size;
234         return memsize;
235 #endif
236
237 default_size:
238         WARNING("Failed to determine available memory; assuming 1 GiB");
239         return (u64)1 << 30;
240 }
241 #endif /* !__WIN32__ */