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 pub fn Set(comptime K: type) type {
  75     return struct {
  76         idx: ?Index = null,
  77
  78         const Self = @This();
  79         pub const Key = K;
  80         pub const Index = u64;
  81         pub const View = SetView(K);
  82
  83         fn open_dbi(txn: lmdb.Txn) !lmdb.Dbi {
  84             return try txn.dbi("SetList");
  85         }
  86         pub fn init(txn: lmdb.Txn) !Self {
  87             const head = View.Head{};
  88             const dbi = try open_dbi(txn);
  89             const idx = try Prng.gen(dbi, Index);
  90             try dbi.put(idx, head);
  91             return .{ .idx = idx };
  92         }
  93         pub fn open(self: Self, txn: lmdb.Txn) !View {
  94             // create new head
  95             if (self.idx == null) {
  96                 return error.NotInitialized;
  97             }
  98             // get head from dbi
  99             const dbi = try open_dbi(txn);
 100             const head = try dbi.get(self.idx.?, View.Head);
 101             return .{
 102                 .dbi = dbi,
 103                 .idx = self.idx.?,
 104                 .head = head,
 105             };
 106         }
 107     };
 108 }
 109
 110 pub fn SetView(comptime K: type) type {
 111     return struct {
 112         const Self = @This();
 113         const ItemIndex = struct { Set(K).Index, K };
 114
 115         pub const Head = struct {
 116             len: usize = 0,
 117             first: ?K = null,
 118             last: ?K = null,
 119         };
 120         pub const Item = struct {
 121             next: ?K = null,
 122             prev: ?K = null,
 123         };
 124
 125         dbi: lmdb.Dbi,
 126         idx: Set(K).Index,
 127         head: Head,
 128
 129         fn item_idx(self: Self, k: K) ItemIndex {
 130             return .{ self.idx, k };
 131         }
 132         fn item_get(self: Self, k: K) !Item {
 133             return try self.dbi.get(self.item_idx(k), Item);
 134         }
 135         fn item_put(self: Self, k: K, item: Item) !void {
 136             try self.dbi.put(self.item_idx(k), item);
 137         }
 138         fn head_update(self: Self) !void {
 139             try self.dbi.put(self.idx, self.head);
 140         }
 141         pub fn append(self: *Self, k: K) !void {
 142             if (self.head.len == 0) {
 143                 const item = Item{};
 144                 try self.item_put(k, item);
 145
 146                 self.head.len = 1;
 147                 self.head.first = k;
 148                 self.head.last = k;
 149                 try self.head_update();
 150             } else {
 151                 const prev_idx = self.head.last.?;
 152                 var prev = try self.item_get(prev_idx);
 153
 154                 const item = Item{ .prev = prev_idx };
 155                 try self.item_put(k, item);
 156
 157                 prev.next = k;
 158                 try self.item_put(prev_idx, prev);
 159
 160                 self.head.last = k;
 161                 self.head.len += 1;
 162                 try self.head_update();
 163             }
 164         }
 165         pub fn del(self: *Self, k: K) !void {
 166             const item = try self.item_get(k);
 167
 168             if (item.prev != null) {
 169                 var prev = try self.item_get(item.prev.?);
 170                 prev.next = item.next;
 171                 try self.item_put(item.prev.?, prev);
 172             }
 173
 174             if (item.next != null) {
 175                 var next = try self.item_get(item.next.?);
 176                 next.prev = item.prev;
 177                 try self.item_put(item.next.?, next);
 178             }
 179
 180             if (self.head.first == k) self.head.first = item.next;
 181             if (self.head.last == k) self.head.last = item.prev;
 182             self.head.len -= 1;
 183             try self.head_update();
 184
 185             try self.dbi.del(self.item_idx(k));
 186         }
 187         pub fn clear(self: *Self) !void {
 188             var it = self.iterator();
 189             while (it.next()) |i| {
 190                 try self.del(i.key);
 191             }
 192         }
 193         pub fn has(self: Self, k: K) !bool {
 194             return self.dbi.has(self.item_idx(k));
 195         }
 196         pub fn len(self: Self) usize {
 197             return self.head.len;
 198         }
 199         pub const Iterator = struct {
 200             sv: SetView(K),
 201             idx: ?K,
 202             dir: enum { Forward, Backward },
 203
 204             pub fn next(self: *Iterator) ?struct { key: K } {
 205                 if (self.idx != null) {
 206                     const k = self.idx.?;
 207                     const item = self.sv.item_get(k) catch return null;
 208                     self.idx = switch (self.dir) {
 209                         .Forward => item.next,
 210                         .Backward => item.prev,
 211                     };
 212                     return .{ .key = k };
 213                 } else {
 214                     return null;
 215                 }
 216             }
 217         };
 218         pub fn iterator(self: Self) Iterator {
 219             return .{
 220                 .sv = self,
 221                 .idx = self.head.first,
 222                 .dir = .Forward,
 223             };
 224         }
 225         pub fn reverse_iterator(self: Self) Iterator {
 226             return .{
 227                 .sv = self,
 228                 .idx = self.head.last,
 229                 .dir = .Backward,
 230             };
 231         }
 232     };
 233 }
 234 pub fn List(comptime V: type) type {
 235     return struct {
 236         idx: ?Index = null,
 237
 238         const Self = @This();
 239         pub const Index = u64;
 240         pub const Key = u64;
 241         pub const Val = V;
 242         pub const View = ListView(V);
 243
 244         fn open_dbi(txn: lmdb.Txn) !lmdb.Dbi {
 245             return try txn.dbi("SetList");
 246         }
 247         pub fn init(txn: lmdb.Txn) !Self {
 248             const head = View.Head{};
 249             const dbi = try open_dbi(txn);
 250             const idx = try Prng.gen(dbi, Index);
 251             try dbi.put(idx, head);
 252             return .{ .idx = idx };
 253         }
 254         pub fn open(self: Self, txn: lmdb.Txn) !View {
 255             // create new head
 256             if (self.idx == null) {
 257                 return error.NotInitialized;
 258             }
 259             // get head from dbi
 260             const dbi = try open_dbi(txn);
 261             const head = try dbi.get(self.idx.?, View.Head);
 262             return .{
 263                 .dbi = dbi,
 264                 .idx = self.idx.?,
 265                 .head = head,
 266             };
 267         }
 268     };
 269 }
 270
 271 pub fn ListView(comptime V: type) type {
 272     return struct {
 273         const Self = @This();
 274         const K = u64;
 275         const ItemIndex = struct { List(V).Index, K };
 276
 277         pub const Head = struct {
 278             len: usize = 0,
 279             first: ?K = null,
 280             last: ?K = null,
 281         };
 282         pub const Item = struct {
 283             next: ?K = null,
 284             prev: ?K = null,
 285             data: V,
 286         };
 287
 288         dbi: lmdb.Dbi,
 289         idx: List(V).Index,
 290         head: Head,
 291
 292         fn item_idx(self: Self, k: K) ItemIndex {
 293             return .{ self.idx, k };
 294         }
 295         fn item_get(self: Self, k: K) !Item {
 296             return try self.dbi.get(self.item_idx(k), Item);
 297         }
 298         fn item_put(self: Self, k: K, item: Item) !void {
 299             try self.dbi.put(self.item_idx(k), item);
 300         }
 301         fn head_update(self: Self) !void {
 302             try self.dbi.put(self.idx, self.head);
 303         }
 304         fn gen(self: Self) !K {
 305             // TODO: limit loop
 306             while (true) {
 307                 const k = try Prng.gen(self.dbi, K);
 308                 if (!try self.dbi.has(self.item_idx(k))) {
 309                     return k;
 310                 }
 311             }
 312         }
 313         pub fn append(self: *Self, v: V) !K {
 314             if (self.head.len == 0) {
 315                 const k = try self.gen();
 316                 const item = Item{ .data = v };
 317                 try self.item_put(k, item);
 318
 319                 self.head.len = 1;
 320                 self.head.first = k;
 321                 self.head.last = k;
 322                 try self.head_update();
 323
 324                 return k;
 325             } else {
 326                 const prev_idx = self.head.last.?;
 327                 var prev = try self.item_get(prev_idx);
 328
 329                 const k = try self.gen();
 330                 const item = Item{ .prev = prev_idx, .data = v };
 331                 try self.item_put(k, item);
 332
 333                 prev.next = k;
 334                 try self.item_put(prev_idx, prev);
 335
 336                 self.head.last = k;
 337                 self.head.len += 1;
 338                 try self.head_update();
 339
 340                 return k;
 341             }
 342         }
 343         pub fn get(self: Self, k: K) !V {
 344             const item = try self.item_get(k);
 345             return item.data;
 346         }
 347         pub fn del(self: *Self, k: K) !void {
 348             const item = try self.item_get(k);
 349
 350             if (item.prev != null) {
 351                 var prev = try self.item_get(item.prev.?);
 352                 prev.next = item.next;
 353                 try self.item_put(item.prev.?, prev);
 354             }
 355
 356             if (item.next != null) {
 357                 var next = try self.item_get(item.next.?);
 358                 next.prev = item.prev;
 359                 try self.item_put(item.next.?, next);
 360             }
 361
 362             if (self.head.first == k) self.head.first = item.next;
 363             if (self.head.last == k) self.head.last = item.prev;
 364             self.head.len -= 1;
 365             try self.head_update();
 366
 367             try self.dbi.del(self.item_idx(k));
 368         }
 369         pub fn clear(self: *Self) !void {
 370             var it = self.iterator();
 371             while (it.next()) |kv| {
 372                 try self.del(kv.key);
 373             }
 374         }
 375         pub fn len(self: Self) usize {
 376             return self.head.len;
 377         }
 378         pub const Iterator = struct {
 379             lv: ListView(V),
 380             idx: ?K,
 381             dir: enum { Forward, Backward },
 382
 383             pub fn next(self: *Iterator) ?struct { key: K, val: V } {
 384                 if (self.idx != null) {
 385                     const k = self.idx.?;
 386                     const item = self.lv.item_get(k) catch return null;
 387                     self.idx = switch (self.dir) {
 388                         .Forward => item.next,
 389                         .Backward => item.prev,
 390                     };
 391                     return .{ .key = k, .val = item.data };
 392                 } else {
 393                     return null;
 394                 }
 395             }
 396         };
 397         pub fn iterator(self: Self) Iterator {
 398             return .{
 399                 .lv = self,
 400                 .idx = self.head.first,
 401                 .dir = .Forward,
 402             };
 403         }
 404         pub fn reverse_iterator(self: Self) Iterator {
 405             return .{
 406                 .lv = self,
 407                 .idx = self.head.last,
 408                 .dir = .Backward,
 409             };
 410         }
 411     };
 412 }
 413
 414 pub fn SetList(comptime K: type, comptime V: type) type {
 415     return struct {
 416         const Self = @This();
 417         pub const Index = u64;
 418         pub const Key = K;
 419         pub const Val = V;
 420         pub const View = SetListView(K, V);
 421
 422         idx: ?Index = null,
 423
 424         fn open_dbi(txn: lmdb.Txn) !lmdb.Dbi {
 425             return try txn.dbi("SetList");
 426         }
 427         pub fn init(txn: lmdb.Txn) !Self {
 428             const head = View.Head{};
 429             const dbi = try open_dbi(txn);
 430             const idx = try Prng.gen(dbi, Index);
 431             try dbi.put(idx, head);
 432             return .{ .idx = idx };
 433         }
 434         pub fn open(self: Self, txn: lmdb.Txn) !View {
 435             // create new head
 436             if (self.idx == null) {
 437                 return error.NotInitialized;
 438             }
 439             // get head from dbi
 440             const dbi = try open_dbi(txn);
 441             const head = try dbi.get(self.idx.?, View.Head);
 442             return .{
 443                 .dbi = dbi,
 444                 .idx = self.idx.?,
 445                 .head = head,
 446             };
 447         }
 448     };
 449 }
 450
 451 pub fn SetListView(comptime K: type, comptime V: type) type {
 452     return struct {
 453         const Self = @This();
 454         const ItemIndex = struct { SetList(K, V).Index, K };
 455
 456         pub const Head = struct {
 457             len: usize = 0,
 458             first: ?K = null,
 459             last: ?K = null,
 460         };
 461         pub const Item = struct {
 462             next: ?K = null,
 463             prev: ?K = null,
 464             data: V,
 465         };
 466
 467         dbi: lmdb.Dbi,
 468         idx: SetList(K, V).Index,
 469         head: Head,
 470
 471         fn item_idx(self: Self, k: K) ItemIndex {
 472             return .{ self.idx, k };
 473         }
 474         fn item_get(self: Self, k: K) !Item {
 475             return try self.dbi.get(self.item_idx(k), Item);
 476         }
 477         fn item_put(self: Self, k: K, item: Item) !void {
 478             try self.dbi.put(self.item_idx(k), item);
 479         }
 480         fn head_update(self: Self) !void {
 481             try self.dbi.put(self.idx, self.head);
 482         }
 483         pub fn append(self: *Self, k: K, v: V) !void {
 484             if (self.head.len == 0) {
 485                 const item = Item{ .data = v };
 486                 try self.item_put(k, item);
 487
 488                 self.head.len = 1;
 489                 self.head.first = k;
 490                 self.head.last = k;
 491                 try self.head_update();
 492             } else {
 493                 const prev_idx = self.head.last.?;
 494                 var prev = try self.item_get(prev_idx);
 495
 496                 const item = Item{ .prev = prev_idx, .data = v };
 497                 try self.item_put(k, item);
 498
 499                 prev.next = k;
 500                 try self.item_put(prev_idx, prev);
 501
 502                 self.head.last = k;
 503                 self.head.len += 1;
 504                 try self.head_update();
 505             }
 506         }
 507         pub fn get(self: Self, k: K) !V {
 508             const item = try self.item_get(k);
 509             return item.data;
 510         }
 511         pub fn del(self: *Self, k: K) !void {
 512             const item = try self.item_get(k);
 513
 514             if (item.prev != null) {
 515                 var prev = try self.item_get(item.prev.?);
 516                 prev.next = item.next;
 517                 try self.item_put(item.prev.?, prev);
 518             }
 519
 520             if (item.next != null) {
 521                 var next = try self.item_get(item.next.?);
 522                 next.prev = item.prev;
 523                 try self.item_put(item.next.?, next);
 524             }
 525
 526             if (self.head.first == k) self.head.first = item.next;
 527             if (self.head.last == k) self.head.last = item.prev;
 528             self.head.len -= 1;
 529             try self.head_update();
 530
 531             try self.dbi.del(self.item_idx(k));
 532         }
 533         pub fn clear(self: *Self) !void {
 534             var it = self.iterator();
 535             while (it.next()) |kv| {
 536                 try self.del(kv.key);
 537             }
 538         }
 539         pub fn has(self: Self, k: K) !bool {
 540             return self.dbi.has(self.item_idx(k));
 541         }
 542         pub fn len(self: Self) usize {
 543             return self.head.len;
 544         }
 545         pub const Iterator = struct {
 546             slv: SetListView(K, V),
 547             idx: ?K,
 548             dir: enum { Forward, Backward },
 549
 550             pub fn next(self: *Iterator) ?struct { key: K, val: V } {
 551                 if (self.idx != null) {
 552                     const k = self.idx.?;
 553                     const item = self.slv.item_get(k) catch return null;
 554                     self.idx = switch (self.dir) {
 555                         .Forward => item.next,
 556                         .Backward => item.prev,
 557                     };
 558                     return .{ .key = k, .val = item.data };
 559                 } else {
 560                     return null;
 561                 }
 562             }
 563         };
 564         pub fn iterator(self: Self) Iterator {
 565             return .{
 566                 .slv = self,
 567                 .idx = self.head.first,
 568                 .dir = .Forward,
 569             };
 570         }
 571         pub fn reverse_iterator(self: Self) Iterator {
 572             return .{
 573                 .slv = self,
 574                 .idx = self.head.last,
 575                 .dir = .Backward,
 576             };
 577         }
 578     };
 579 }
 580
 581 const DB_SIZE = 1024 * 1024 * 1;
 582
 583 test "db" {
 584     const env = try lmdb.Env.open("db", DB_SIZE);
 585     defer env.close();
 586
 587     const txn = try env.txn();
 588     defer txn.commit() catch {};
 589
 590     var db = try Db(u32, u32).init(txn, "123");
 591     var n: u32 = 456;
 592     if (try db.has(123)) {
 593         n = try db.get(123);
 594         n += 1;
 595     }
 596     try db.put(123, n);
 597     std.debug.print("n: {}\n", .{n});
 598 }
 599
 600 // test "list" {
 601 //     const env = try lmdb.Env.open("db", DB_SIZE);
 602 //     defer env.close();
 603
 604 //     const txn = try env.txn();
 605 //     defer txn.commit();
 606
 607 //     const db = List.init(txn, "b", u32);
 608 // }
 609
 610 test "set" {
 611     var env = try lmdb.Env.open("db", 1024 * 1024 * 1);
 612     // env.sync();
 613     defer env.close();
 614
 615     var txn = try env.txn();
 616     defer txn.commit() catch {};
 617
 618     var dbi = try txn.dbi("abc");
 619
 620     const A = struct {
 621         ml: Set(usize),
 622     };
 623
 624     var a: A = undefined;
 625     const a_idx: u64 = 27;
 626     if (try dbi.has(a_idx)) {
 627         a = try dbi.get(a_idx, A);
 628     } else {
 629         a = A{ .ml = try Set(usize).init(txn) };
 630         try dbi.put(a_idx, a);
 631     }
 632
 633     var ml = try a.ml.open(txn);
 634
 635     const len = ml.len();
 636     std.debug.print("{}\n", .{len});
 637     try ml.append(len);
 638     std.debug.print("{}\n", .{try ml.has(len)});
 639     var it = ml.iterator();
 640     while (it.next()) |i| {
 641         std.debug.print("{}\n", .{i});
 642     }
 643 }
 644
 645 test "list" {
 646     var env = try lmdb.Env.open("db", 1024 * 1024 * 1);
 647     // env.sync();
 648     defer env.close();
 649
 650     var txn = try env.txn();
 651     defer txn.commit() catch {};
 652
 653     var dbi = try txn.dbi("def");
 654
 655     const A = struct {
 656         ml: List(usize),
 657     };
 658
 659     var a: A = undefined;
 660     const a_idx: u64 = 27;
 661     if (try dbi.has(a_idx)) {
 662         a = try dbi.get(a_idx, A);
 663     } else {
 664         a = A{ .ml = try List(usize).init(txn) };
 665         try dbi.put(a_idx, a);
 666     }
 667
 668     var ml = try a.ml.open(txn);
 669
 670     const len = ml.len();
 671     std.debug.print("{}\n", .{len});
 672     const newest = try ml.append(len * 10);
 673     std.debug.print("{}: {}\n", .{ newest, try ml.get(newest) });
 674     var it = ml.iterator();
 675     while (it.next()) |i| {
 676         std.debug.print("{}: {}\n", .{ i.key, i.val });
 677     }
 678 }
 679
 680 test "setlist" {
 681     var env = try lmdb.Env.open("db", 1024 * 1024 * 1);
 682     // env.sync();
 683     defer env.close();
 684
 685     var txn = try env.txn();
 686     defer txn.commit() catch {};
 687
 688     var dbi = try txn.dbi("ghi");
 689
 690     const A = struct {
 691         ml: SetList(usize, usize),
 692     };
 693
 694     var a: A = undefined;
 695     const a_idx: u64 = 27;
 696     if (try dbi.has(a_idx)) {
 697         a = try dbi.get(a_idx, A);
 698     } else {
 699         a = A{ .ml = try SetList(usize, usize).init(txn) };
 700         try dbi.put(a_idx, a);
 701     }
 702
 703     var ml = try a.ml.open(txn);
 704
 705     const len = ml.len();
 706     std.debug.print("{}\n", .{len});
 707     try ml.append(len, len * 10);
 708     std.debug.print("{}\n", .{try ml.get(len)});
 709     var it = ml.iterator();
 710     while (it.next()) |i| {
 711         std.debug.print("{}: {}\n", .{ i.key, i.val });
 712     }
 713 }