Coverage for src/evutils/io/_native_core.py: 81%

1"""Low-level ctypes binding and common structs for the native library. 

2""" 

3from __future__ import annotations 

4 

5import ctypes 

6import os 

7import sys 

8from ctypes import ( 

9 POINTER, Structure, byref, cast as c_cast, c_char_p, c_char, 

10 c_int, c_size_t, c_uint8, c_uint16, c_uint64, c_void_p, 

11) 

12from pathlib import Path 

13from typing import cast, Any, Callable 

14 

15import numpy as np 

16from ..types import EventArray, TriggerArray 

17 

18__all__ = [ 

19 "NativeError", "lib", "register_bindings", 

20 "EventBufferSOA", "TriggerBufferSOA", 

21 "ParserResult", "EventSoABuffers", "TriggerSoABuffers", 

22 "EVENT_DTYPE", "TRIGGER_DTYPE", 

23 "events_view", "triggers_view", 

24 "parse_step", "decode_all_soa", 

25 "EVUTILS_PARSE_OK", "EVUTILS_PARSE_INPUT_EMPTY", "EVUTILS_PARSE_OUTPUT_FULL", 

26 "EVUTILS_PARSE_ERROR", "EVUTILS_PARSE_INCOMPLETE", 

27 "_T_DTYPE", "_X_DTYPE", "_Y_DTYPE", "_P_DTYPE", "_ID_DTYPE" 

28] 

29 

30class NativeError(RuntimeError): 

31 pass 

32 

33EVENT_DTYPE = np.dtype({"names": ["t", "x", "y", "p"], "formats": ["<u4", "<u2", "<u2", "u1"], "offsets": [0, 4, 6, 8], "itemsize": 12}) 

34TRIGGER_DTYPE = np.dtype({"names": ["t", "id", "p"], "formats": ["<u4", "u1", "u1"], "offsets": [0, 4, 5], "itemsize": 8}) 

35 

36_T_DTYPE = np.uint64 

37_X_DTYPE = np.uint16 

38_Y_DTYPE = np.uint16 

39_P_DTYPE = np.uint8 

40_ID_DTYPE = np.uint8 

41 

42class EventBufferSOA(Structure): 

43 _fields_ = [ 

44 ("t", POINTER(c_uint64)), ("x", POINTER(c_uint16)), 

45 ("y", POINTER(c_uint16)), ("p", POINTER(c_uint8)), 

46 ("capacity", c_size_t), ("size", c_size_t), 

47 ] 

48 

49class TriggerBufferSOA(Structure): 

50 _fields_ = [ 

51 ("t", POINTER(c_uint64)), ("id", POINTER(c_uint8)), ("p", POINTER(c_uint8)), 

52 ("capacity", c_size_t), ("size", c_size_t), 

53 ] 

54 

55EVUTILS_PARSE_OK = 0 

56EVUTILS_PARSE_INPUT_EMPTY = 1 

57EVUTILS_PARSE_OUTPUT_FULL = 2 

58EVUTILS_PARSE_ERROR = 3 

59EVUTILS_PARSE_INCOMPLETE = 4 

60 

61class ParserResult(Structure): 

62 _fields_ = [("current", POINTER(c_uint16)), ("status", c_int)] 

63 

64def _candidate_filenames() -> list[str]: 

65 base = "evutils_native" 

66 if sys.platform.startswith("win"): return [f"{base}.dll", f"lib{base}.dll"] 

67 if sys.platform == "darwin": return [f"lib{base}.dylib", f"{base}.dylib"] 

68 return [f"lib{base}.so", f"{base}.so"] 

69 

70def _search_roots() -> list[Path]: 

71 here = Path(__file__).resolve().parent 

72 roots = [here] 

73 for parent in (here, *here.parents[:5]): 

74 if (parent / "pyproject.toml").exists(): 

75 roots.append(parent / "build") 

76 break 

77 roots.append(Path.cwd() / "build") 

78 return roots 

79 

80def _find_library() -> str: 

81 override = os.environ.get("EVUTILS_NATIVE_LIB") 

82 if override: return override 

83 names = set(_candidate_filenames()) 

84 for root in _search_roots(): 

85 if not root.exists(): continue 

86 for name in names: 

87 p = root / name 

88 if p.is_file(): return str(p) 

89 for p in root.glob("**/*evutils_native*"): 

90 if p.is_file() and p.name in names: return str(p) 

91 raise NativeError("Could not find the evutils native library.") 

92 

93_BINDINGS: list[Callable[[ctypes.CDLL], None]] = [] 

94 

95def register_bindings(binder: Callable[[ctypes.CDLL], None]) -> None: 

96 _BINDINGS.append(binder) 

97 # If already loaded, bind immediately 

98 if _LIB is not None: 

99 binder(_LIB) 

100 

101def _bind(handle: ctypes.CDLL) -> ctypes.CDLL: 

102 handle.evutils_version.argtypes = [] 

103 handle.evutils_version.restype = c_char_p 

104 for binder in _BINDINGS: 

105 binder(handle) 

106 return handle 

107 

108_LIB: ctypes.CDLL | None = None 

109 

110def lib() -> ctypes.CDLL: 

111 global _LIB 

112 if _LIB is None: 

113 try: 

114 handle = ctypes.CDLL(_find_library()) 

115 except OSError as exc: 

116 raise NativeError(f"Failed to load evutils native library: {exc}") from exc 

117 _LIB = _bind(handle) 

118 return _LIB 

119 

120class EventSoABuffers: 

121 __slots__ = ("capacity", "t", "x", "y", "p", "c") 

122 def __init__(self, capacity: int): 

123 self.capacity = int(capacity) 

124 self.t = np.empty(self.capacity, dtype=_T_DTYPE) 

125 self.x = np.empty(self.capacity, dtype=_X_DTYPE) 

126 self.y = np.empty(self.capacity, dtype=_Y_DTYPE) 

127 self.p = np.empty(self.capacity, dtype=_P_DTYPE) 

128 self.c = EventBufferSOA() 

129 self.c.t = self.t.ctypes.data_as(POINTER(c_uint64)) 

130 self.c.x = self.x.ctypes.data_as(POINTER(c_uint16)) 

131 self.c.y = self.y.ctypes.data_as(POINTER(c_uint16)) 

132 self.c.p = self.p.ctypes.data_as(POINTER(c_uint8)) 

133 self.c.capacity = self.capacity 

134 self.c.size = 0 

135 @property 

136 def size(self) -> int: return int(self.c.size) 

137 def reset(self) -> None: self.c.size = 0 

138 def grow(self, new_capacity: int) -> None: 

139 new_capacity = int(new_capacity) 

140 if new_capacity <= self.capacity: return 

141 n = self.size 

142 for name, dtype, field, ptr in ( 

143 ("t", _T_DTYPE, "t", c_uint64), ("x", _X_DTYPE, "x", c_uint16), 

144 ("y", _Y_DTYPE, "y", c_uint16), ("p", _P_DTYPE, "p", c_uint8), 

145 ): 

146 grown = np.empty(new_capacity, dtype=dtype) 

147 grown[:n] = getattr(self, name)[:n] 

148 setattr(self, name, grown) 

149 setattr(self.c, field, grown.ctypes.data_as(POINTER(ptr))) 

150 self.capacity = new_capacity 

151 self.c.capacity = new_capacity 

152 def view(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: 

153 n = self.size 

154 return self.t[:n], self.x[:n], self.y[:n], self.p[:n] 

155 

156class TriggerSoABuffers: 

157 __slots__ = ("capacity", "t", "id", "p", "c") 

158 def __init__(self, capacity: int): 

159 self.capacity = int(capacity) 

160 self.t = np.empty(self.capacity, dtype=_T_DTYPE) 

161 self.id = np.empty(self.capacity, dtype=_ID_DTYPE) 

162 self.p = np.empty(self.capacity, dtype=_P_DTYPE) 

163 self.c = TriggerBufferSOA() 

164 self.c.t = self.t.ctypes.data_as(POINTER(c_uint64)) 

165 self.c.id = self.id.ctypes.data_as(POINTER(c_uint8)) 

166 self.c.p = self.p.ctypes.data_as(POINTER(c_uint8)) 

167 self.c.capacity = self.capacity 

168 self.c.size = 0 

169 @property 

170 def size(self) -> int: return int(self.c.size) 

171 def reset(self) -> None: self.c.size = 0 

172 def grow(self, new_capacity: int) -> None: 

173 new_capacity = int(new_capacity) 

174 if new_capacity <= self.capacity: return 

175 n = self.size 

176 for name, dtype, ptr in ( 

177 ("t", _T_DTYPE, c_uint64), ("id", _ID_DTYPE, c_uint8), ("p", _P_DTYPE, c_uint8), 

178 ): 

179 grown = np.empty(new_capacity, dtype=dtype) 

180 grown[:n] = getattr(self, name)[:n] 

181 setattr(self, name, grown) 

182 setattr(self.c, name, grown.ctypes.data_as(POINTER(ptr))) 

183 self.capacity = new_capacity 

184 self.c.capacity = new_capacity 

185 def view(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: 

186 n = self.size 

187 return self.t[:n], self.id[:n], self.p[:n] 

188 

189def events_view(ev: EventSoABuffers) -> EventArray: 

190 n = ev.size 

191 return EventArray(ev.t[:n].view(np.int64), ev.x[:n], ev.y[:n], ev.p[:n]) 

192 

193def triggers_view(tr: TriggerSoABuffers) -> TriggerArray: 

194 n = tr.size 

195 return TriggerArray(tr.t[:n].view(np.int64), tr.p[:n], tr.id[:n]) 

196 

197def parse_step(words: Any, offset: int, make_input: Any, parser: Any, events: Any, triggers: Any, *, tail_pad: int = 0, word_dtype: Any = None) -> tuple[int, int]: 

198 n_words = len(words) 

199 if offset >= n_words: return 0, n_words 

200 before = events.size 

201 inp = make_input(words[offset:]) 

202 res = parser.parse_chunk_soa(inp, events, triggers) 

203 if res.status == EVUTILS_PARSE_ERROR: raise RuntimeError(f"parse error near word {offset}") 

204 consumed = inp.consumed(res) 

205 if consumed == 0: 

206 if tail_pad and word_dtype is not None: 

207 tail = words[offset:] 

208 if len(tail): 

209 scratch = np.zeros(len(tail) + tail_pad, dtype=word_dtype) 

210 scratch[: len(tail)] = tail 

211 parser.parse_chunk_soa(make_input(scratch), events, triggers) 

212 return events.size - before, n_words 

213 return events.size - before, offset + consumed 

214 

215def decode_all_soa(words: Any, start_offset: int, make_input: Any, parser: Any, *, est_events_per_word: float = 1.0, tail_pad: int = 0, word_dtype: Any = None, trigger_cap: int = 1 << 12) -> tuple[EventArray, int]: 

216 n_words = len(words) if words is not None else 0 

217 if n_words == 0 or start_offset >= n_words: return EventArray.empty(), n_words 

218 remaining = n_words - start_offset 

219 cap = int(remaining * est_events_per_word) + 1024 

220 ev = EventSoABuffers(cap) 

221 tr = TriggerSoABuffers(max(trigger_cap, 1)) 

222 offset = start_offset 

223 while offset < n_words: 

224 if ev.capacity - ev.size < 128: ev.grow(int(ev.capacity * 1.5) + (1 << 16)) 

225 if tr.capacity - tr.size < 64: tr.grow(tr.capacity * 2 + 64) 

226 inp = make_input(words[offset:]) 

227 res = parser.parse_chunk_soa(inp, ev, tr) 

228 if res.status == EVUTILS_PARSE_ERROR: raise RuntimeError(f"parse error near word {offset}") 

229 consumed = inp.consumed(res) 

230 if consumed == 0: 

231 if tail_pad and word_dtype is not None: 

232 tail = words[offset:] 

233 if len(tail): 

234 if ev.capacity - ev.size < len(tail) + 128: ev.grow(ev.size + len(tail) + (1 << 16)) 

235 scratch = np.zeros(len(tail) + tail_pad, dtype=word_dtype) 

236 scratch[: len(tail)] = tail 

237 parser.parse_chunk_soa(make_input(scratch), ev, tr) 

238 offset = n_words 

239 break 

240 offset += consumed 

241 n = ev.size 

242 if n == 0: return EventArray.empty(), offset 

243 out = EventArray(ev.t[:n].view(np.int64), ev.x[:n], ev.y[:n], ev.p[:n]) 

244 return out, offset