gitweb.ps.run Git - ziglmdb/blob - src/db.zig

   1 const std = @import("std");
   2 const lmdb = @import("lmdb");
   3
   4 const PRNG_SEED = 0;
   5
   6 pub fn Db(comptime K: type, comptime V: type) type {
   7     return struct {
   8         const Self = @This();
   9
  10         dbi: lmdb.Dbi,
  11
  12         pub fn init(txn: lmdb.Txn, name: [:0]const u8) !Self {
  13             return .{
  14                 .dbi = try txn.dbi(name),
  15             };
  16         }
  17         pub fn put(self: Self, k: K, v: V) !void {
  18             try self.dbi.put(k, v);
  19         }
  20         pub fn get(self: Self, k: K) !V {
  21             return try self.dbi.get(k, V);
  22         }
  23         pub fn del(self: Self, k: K) !void {
  24             try self.dbi.del(k);
  25         }
  26         pub fn has(self: Self, k: K) !bool {
  27             return try self.dbi.has(k);
  28         }
  29         pub const Iterator = struct {
  30             cursor: lmdb.Cursor,
  31             k: ?K,
  32             v: ?V,
  33
  34             pub fn next(self: *Iterator) ?struct { key: K, val: V } {
  35                 if (self.k != null and self.v != null) {
  36                     const result = .{ .key = self.k.?, .val = self.v.? };
  37
  38                     var k = self.k.?;
  39                     self.v = self.cursor.get(&k, V, .Next) catch return null;
  40                     if (self.v != null) {
  41                         self.k = k;
  42                     }
  43                     return result;
  44                 } else {
  45                     return null;
  46                 }
  47             }
  48         };
  49         pub fn iterator(self: Self) !Iterator {
  50             var cursor = try self.dbi.cursor();
  51
  52             var k: K = undefined;
  53             const v = try cursor.get(&k, V, .First);
  54             return .{ .cursor = cursor, .k = k, .v = v };
  55         }
  56     };
  57 }
  58
  59 pub const Prng = struct {
  60     var prng = std.Random.DefaultPrng.init(PRNG_SEED);
  61
  62     pub fn gen(dbi: lmdb.Dbi, comptime T: type) !T {
  63         var buf: [@sizeOf(T)]u8 = undefined;
  64         // TODO: limit loop
  65         while (true) {
  66             prng.fill(&buf);
  67             const t = std.mem.bytesToValue(T, &buf);
  68             if (!try dbi.has(t)) {
  69                 return t;
  70             }
  71         }
  72     }
  73 };
  74
  75 fn SetListBase(comptime K: type, comptime V: type) type {
  76     return struct {
  77         const Self = @This();
  78         pub const Index = u64;
  79         pub const Key = K;
  80         pub const Val = V;
  81         pub const View = SetListViewBase(K, V);
  82
  83         idx: ?Index = null,
  84
  85         fn open_dbi(txn: lmdb.Txn) !lmdb.Dbi {
  86             return try txn.dbi("SetList");
  87         }
  88         pub fn init(txn: lmdb.Txn) !Self {
  89             const head = View.Head{};
  90             const dbi = try open_dbi(txn);
  91             const idx = try Prng.gen(dbi, Index);
  92             try dbi.put(idx, head);
  93             return .{ .idx = idx };
  94         }
  95         pub fn open(self: Self, txn: lmdb.Txn) !View {
  96             // create new head
  97             if (self.idx == null) {
  98                 return error.NotInitialized;
  99             }
 100             // get head from dbi
 101             const dbi = try open_dbi(txn);
 102             const head = try dbi.get(self.idx.?, View.Head);
 103             return .{
 104                 .dbi = dbi,
 105                 .idx = self.idx.?,
 106                 .head = head,
 107             };
 108         }
 109     };
 110 }
 111
 112 fn SetListViewBase(comptime K: type, comptime V: type) type {
 113     return struct {
 114         const Self = @This();
 115         pub const ItemIndex = struct { SetListBase(K, V).Index, Key };
 116         pub const Key = K;
 117         pub const Val = V;
 118
 119         pub const Head = struct {
 120             len: usize = 0,
 121             first: ?K = null,
 122             last: ?K = null,
 123         };
 124         pub const Item = struct {
 125             next: ?K = null,
 126             prev: ?K = null,
 127             data: V,
 128         };
 129
 130         dbi: lmdb.Dbi,
 131         idx: SetListBase(K, V).Index,
 132         head: Head,
 133
 134         fn item_idx(self: Self, k: K) ItemIndex {
 135             return .{ self.idx, k };
 136         }
 137         fn item_get(self: Self, k: K) !Item {
 138             return try self.dbi.get(self.item_idx(k), Item);
 139         }
 140         fn item_put(self: Self, k: K, item: Item) !void {
 141             try self.dbi.put(self.item_idx(k), item);
 142         }
 143         fn head_update(self: Self) !void {
 144             try self.dbi.put(self.idx, self.head);
 145         }
 146         pub fn del(self: *Self, k: K) !void {
 147             const item = try self.item_get(k);
 148
 149             if (item.prev != null) {
 150                 var prev = try self.item_get(item.prev.?);
 151                 prev.next = item.next;
 152                 try self.item_put(item.prev.?, prev);
 153             }
 154
 155             if (item.next != null) {
 156                 var next = try self.item_get(item.next.?);
 157                 next.prev = item.prev;
 158                 try self.item_put(item.next.?, next);
 159             }
 160
 161             if (self.head.first == k) self.head.first = item.next;
 162             if (self.head.last == k) self.head.last = item.prev;
 163             self.head.len -= 1;
 164             try self.head_update();
 165
 166             try self.dbi.del(self.item_idx(k));
 167         }
 168         pub fn clear(self: *Self) !void {
 169             var it = self.iterator();
 170             while (it.next()) |kv| {
 171                 try self.del(kv.key);
 172             }
 173         }
 174         pub fn len(self: Self) usize {
 175             return self.head.len;
 176         }
 177         pub fn append(self: *Self, key: Key, val: Val) !void {
 178             if (self.head.len == 0) {
 179                 const item = Item{ .data = val };
 180                 try self.item_put(key, item);
 181
 182                 self.head.len = 1;
 183                 self.head.first = key;
 184                 self.head.last = key;
 185                 try self.head_update();
 186             } else {
 187                 const prev_idx = self.head.last.?;
 188                 var prev = try self.item_get(prev_idx);
 189
 190                 const item = Item{ .prev = prev_idx, .data = val };
 191                 try self.item_put(key, item);
 192
 193                 prev.next = key;
 194                 try self.item_put(prev_idx, prev);
 195
 196                 self.head.last = key;
 197                 self.head.len += 1;
 198                 try self.head_update();
 199             }
 200         }
 201         pub fn get(self: Self, key: Key) !Val {
 202             const item = try self.item_get(key);
 203             return item.data;
 204         }
 205         pub fn has(self: Self, key: Key) !bool {
 206             return self.dbi.has(self.item_idx(key));
 207         }
 208         pub const Iterator = struct {
 209             pub const Result = ?struct { key: K, val: V };
 210
 211             slv: SetListViewBase(K, V),
 212             idx: ?K,
 213             dir: enum { Forward, Backward },
 214
 215             pub fn next(self: *Iterator) Result {
 216                 if (self.idx != null) {
 217                     const k = self.idx.?;
 218                     const item = self.slv.item_get(k) catch return null;
 219                     self.idx = switch (self.dir) {
 220                         .Forward => item.next,
 221                         .Backward => item.prev,
 222                     };
 223                     return .{ .key = k, .val = item.data };
 224                 } else {
 225                     return null;
 226                 }
 227             }
 228         };
 229         pub fn iterator(self: Self) Iterator {
 230             return .{
 231                 .slv = self,
 232                 .idx = self.head.first,
 233                 .dir = .Forward,
 234             };
 235         }
 236         pub fn reverse_iterator(self: Self) Iterator {
 237             return .{
 238                 .slv = self,
 239                 .idx = self.head.last,
 240                 .dir = .Backward,
 241             };
 242         }
 243     };
 244 }
 245
 246 pub fn Set(comptime K: type) type {
 247     return struct {
 248         pub const Key = K;
 249         pub const Val = void;
 250
 251         pub const Base = SetListBase(Key, Val);
 252         pub const View = struct {
 253             const ViewBase = SetListViewBase(Key, Val);
 254
 255             base: ViewBase,
 256
 257             pub fn del(self: *@This(), key: Key) !void {
 258                 try self.base.del(key);
 259             }
 260             pub fn clear(self: *@This()) !void {
 261                 try self.base.clear();
 262             }
 263             pub fn len(self: @This()) usize {
 264                 return self.base.len();
 265             }
 266             pub fn append(self: *@This(), key: Key) !void {
 267                 try self.base.append(key, {});
 268             }
 269             pub fn has(self: @This(), key: Key) !bool {
 270                 return try self.base.has(key);
 271             }
 272             pub fn iterator(self: @This()) ViewBase.Iterator {
 273                 return self.base.iterator();
 274             }
 275             pub fn reverse_iterator(self: @This()) ViewBase.Iterator {
 276                 return self.base.reverse_iterator();
 277             }
 278         };
 279
 280         base: Base,
 281
 282         pub fn init(txn: lmdb.Txn) !@This() {
 283             return .{ .base = try Base.init(txn) };
 284         }
 285         pub fn open(self: @This(), txn: lmdb.Txn) !View {
 286             return .{ .base = try self.base.open(txn) };
 287         }
 288     };
 289 }
 290
 291 pub fn List(comptime V: type) type {
 292     return struct {
 293         pub const Key = u64;
 294         pub const Val = V;
 295
 296         pub const Base = SetListBase(Key, Val);
 297         pub const View = struct {
 298             const ViewBase = SetListViewBase(Key, Val);
 299
 300             base: ViewBase,
 301
 302             fn gen(self: @This()) !Key {
 303                 // TODO: limit loop
 304                 while (true) {
 305                     const key = try Prng.gen(self.base.dbi, Key);
 306                     if (!try self.base.dbi.has(self.base.item_idx(key))) {
 307                         return key;
 308                     }
 309                 }
 310             }
 311             pub fn del(self: *@This(), key: Key) !void {
 312                 try self.base.del(key);
 313             }
 314             pub fn clear(self: *@This()) !void {
 315                 try self.base.clear();
 316             }
 317             pub fn len(self: @This()) usize {
 318                 return self.base.len();
 319             }
 320             pub fn append(self: *@This(), val: Val) !Key {
 321                 const key = try self.gen();
 322                 try self.base.append(key, val);
 323                 return key;
 324             }
 325             pub fn get(self: @This(), key: Key) !Val {
 326                 return try self.base.get(key);
 327             }
 328             pub fn has(self: @This(), key: Key) !bool {
 329                 return try self.base.has(key);
 330             }
 331             pub fn iterator(self: @This()) ViewBase.Iterator {
 332                 return self.base.iterator();
 333             }
 334             pub fn reverse_iterator(self: @This()) ViewBase.Iterator {
 335                 return self.base.reverse_iterator();
 336             }
 337         };
 338
 339         base: Base,
 340
 341         pub fn init(txn: lmdb.Txn) !@This() {
 342             return .{ .base = try Base.init(txn) };
 343         }
 344         pub fn open(self: @This(), txn: lmdb.Txn) !View {
 345             return .{ .base = try self.base.open(txn) };
 346         }
 347     };
 348 }
 349
 350 pub fn SetList(comptime K: type, comptime V: type) type {
 351     return struct {
 352         pub const Key = K;
 353         pub const Val = V;
 354
 355         pub const Base = SetListBase(Key, Val);
 356         pub const View = struct {
 357             const ViewBase = SetListViewBase(Key, Val);
 358
 359             base: ViewBase,
 360
 361             pub fn del(self: *@This(), key: Key) !void {
 362                 try self.base.del(key);
 363             }
 364             pub fn clear(self: *@This()) !void {
 365                 try self.base.clear();
 366             }
 367             pub fn len(self: @This()) usize {
 368                 return self.base.len();
 369             }
 370             pub fn append(self: *@This(), key: Key, val: Val) !void {
 371                 try self.base.append(key, val);
 372             }
 373             pub fn get(self: @This(), key: Key) !Val {
 374                 return try self.base.get(key);
 375             }
 376             pub fn has(self: @This(), key: Key) !bool {
 377                 return try self.base.has(key);
 378             }
 379             pub fn iterator(self: @This()) ViewBase.Iterator {
 380                 return self.base.iterator();
 381             }
 382             pub fn reverse_iterator(self: @This()) ViewBase.Iterator {
 383                 return self.base.reverse_iterator();
 384             }
 385         };
 386
 387         base: Base,
 388
 389         pub fn init(txn: lmdb.Txn) !@This() {
 390             return .{ .base = try Base.init(txn) };
 391         }
 392         pub fn open(self: @This(), txn: lmdb.Txn) !View {
 393             return .{ .base = try self.base.open(txn) };
 394         }
 395     };
 396 }
 397
 398 const DB_SIZE = 1024 * 1024 * 1;
 399
 400 test "db" {
 401     const env = try lmdb.Env.open("db", DB_SIZE);
 402     defer env.close();
 403
 404     const txn = try env.txn();
 405     defer txn.commit() catch {};
 406
 407     var db = try Db(u32, u32).init(txn, "123");
 408     var n: u32 = 456;
 409     if (try db.has(123)) {
 410         n = try db.get(123);
 411         n += 1;
 412     }
 413     try db.put(123, n);
 414     std.debug.print("n: {}\n", .{n});
 415 }
 416
 417 // test "list" {
 418 //     const env = try lmdb.Env.open("db", DB_SIZE);
 419 //     defer env.close();
 420
 421 //     const txn = try env.txn();
 422 //     defer txn.commit();
 423
 424 //     const db = List.init(txn, "b", u32);
 425 // }
 426
 427 test "set" {
 428     var env = try lmdb.Env.open("db", 1024 * 1024 * 1);
 429     // env.sync();
 430     defer env.close();
 431
 432     var txn = try env.txn();
 433     defer txn.commit() catch {};
 434
 435     var dbi = try txn.dbi("abc");
 436
 437     const A = struct {
 438         ml: Set(usize),
 439     };
 440
 441     var a: A = undefined;
 442     const a_idx: u64 = 27;
 443     if (try dbi.has(a_idx)) {
 444         a = try dbi.get(a_idx, A);
 445     } else {
 446         a = A{ .ml = try Set(usize).init(txn) };
 447         try dbi.put(a_idx, a);
 448     }
 449
 450     var ml = try a.ml.open(txn);
 451
 452     const len = ml.len();
 453     std.debug.print("{}\n", .{len});
 454     try ml.append(len);
 455     std.debug.print("{}\n", .{try ml.has(len)});
 456     var it = ml.iterator();
 457     while (it.next()) |i| {
 458         std.debug.print("{}\n", .{i});
 459     }
 460 }
 461
 462 test "list" {
 463     var env = try lmdb.Env.open("db", 1024 * 1024 * 1);
 464     // env.sync();
 465     defer env.close();
 466
 467     var txn = try env.txn();
 468     defer txn.commit() catch {};
 469
 470     var dbi = try txn.dbi("def");
 471
 472     const A = struct {
 473         ml: List(usize),
 474     };
 475
 476     var a: A = undefined;
 477     const a_idx: u64 = 27;
 478     if (try dbi.has(a_idx)) {
 479         a = try dbi.get(a_idx, A);
 480     } else {
 481         a = A{ .ml = try List(usize).init(txn) };
 482         try dbi.put(a_idx, a);
 483     }
 484
 485     var ml = try a.ml.open(txn);
 486
 487     const len = ml.len();
 488     std.debug.print("{}\n", .{len});
 489     const newest = try ml.append(len * 10);
 490     std.debug.print("{}: {}\n", .{ newest, try ml.get(newest) });
 491     var it = ml.iterator();
 492     while (it.next()) |i| {
 493         std.debug.print("{}: {}\n", .{ i.key, i.val });
 494     }
 495 }
 496
 497 test "setlist" {
 498     var env = try lmdb.Env.open("db", 1024 * 1024 * 1);
 499     // env.sync();
 500     defer env.close();
 501
 502     var txn = try env.txn();
 503     defer txn.commit() catch {};
 504
 505     var dbi = try txn.dbi("ghi");
 506
 507     const A = struct {
 508         ml: SetList(usize, usize),
 509     };
 510
 511     var a: A = undefined;
 512     const a_idx: u64 = 27;
 513     if (try dbi.has(a_idx)) {
 514         a = try dbi.get(a_idx, A);
 515     } else {
 516         a = A{ .ml = try SetList(usize, usize).init(txn) };
 517         try dbi.put(a_idx, a);
 518     }
 519
 520     var ml = try a.ml.open(txn);
 521
 522     const len = ml.len();
 523     std.debug.print("{}\n", .{len});
 524     try ml.append(len, len * 10);
 525     std.debug.print("{}\n", .{try ml.get(len)});
 526     var it = ml.iterator();
 527     while (it.next()) |i| {
 528         std.debug.print("{}: {}\n", .{ i.key, i.val });
 529     }
 530 }