lib/pluck/src/pexpr.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const log = @import("logger.zig");
3 const builtin = @import("builtin");
4 const Allocator = std.mem.Allocator;
5
6 pub const Symbol = []const u8;
7
8 pub const NativeValue = union(enum) {
9 int: i64,
10 float: f64,
11 symbol: Symbol,
12 bool_val: bool,
13
14 pub fn format(
15 self: NativeValue,
16 comptime fmt: []const u8,
17 options: std.fmt.Options,
18 writer: anytype,
19 ) !void {
20 _ = fmt;
21 _ = options;
22 switch (self) {
23 .int => |i| try writer.print("@{d}", .{i}),
24 .float => |f| try writer.print("{d}", .{f}),
25 .symbol => |s| try writer.print("'{s}", .{s}),
26 .bool_val => |b| try writer.print("{}", .{b}),
27 }
28 }
29 };
30
31 pub const CaseOfGuard = struct {
32 constructor: Symbol,
33 args: []const Symbol,
34
35 pub fn format(
36 self: CaseOfGuard,
37 comptime fmt: []const u8,
38 options: std.fmt.Options,
39 writer: anytype,
40 ) !void {
41 _ = fmt;
42 _ = options;
43 try writer.print("{s}", .{self.constructor});
44 for (self.args) |arg| {
45 try writer.print(" {s}", .{arg});
46 }
47 }
48 };
49
50 pub const ConstructorDef = struct {
51 name: Symbol,
52 args: []const Symbol,
53 };
54
55 pub const Head = union(enum) {
56 app: void,
57 abs: struct { var_name: Symbol },
58 var_ref: struct { name: Symbol },
59 defined: struct { name: Symbol },
60
61 const_native: NativeValue,
62
63 case_of: struct { branches: []const CaseOfGuard },
64
65 construct: struct { constructor: Symbol },
66
67 type_def: struct {
68 type_name: Symbol,
69 constructors: []const ConstructorDef,
70 },
71
72 y_combinator: void,
73 flip: void,
74 factor: void,
75 native_eq: void,
76 get_args: void,
77 get_constructor: void,
78 pbool: void,
79 get_config: void,
80 mk_int: void,
81 mk_int_weighted: void,
82 int_dist_eq: void,
83 print_op: void,
84 f_div: void,
85 f_mul: void,
86 f_add: void,
87 f_sub: void,
88 error_op: void,
89
90 pub fn format(
91 self: Head,
92 comptime fmt: []const u8,
93 options: std.fmt.Options,
94 writer: anytype,
95 ) !void {
96 switch (self) {
97 .app => try writer.writeAll("App"),
98 .abs => |a| try writer.print("λ{s}", .{a.var_name}),
99 .var_ref => |v| try writer.print("${s}", .{v.name}),
100 .defined => |d| try writer.print("{s}", .{d.name}),
101 .const_native => |c| try c.format(fmt, options, writer),
102 .case_of => try writer.writeAll("caseof"),
103 .construct => |c| try writer.print("{s}", .{c.constructor}),
104 .type_def => |t| try writer.print("type {s}", .{t.type_name}),
105 .y_combinator => try writer.writeAll("Y"),
106 .flip => try writer.writeAll("flip"),
107 .factor => try writer.writeAll("factor"),
108 .native_eq => try writer.writeAll("native_eq"),
109 .get_args => try writer.writeAll("get_args"),
110 .get_constructor => try writer.writeAll("get_constructor"),
111 .pbool => try writer.writeAll("pbool"),
112 .get_config => try writer.writeAll("get_config"),
113 .mk_int => try writer.writeAll("mk_int"),
114 .mk_int_weighted => try writer.writeAll("mk_int_weighted"),
115 .int_dist_eq => try writer.writeAll("int_dist_eq"),
116 .print_op => try writer.writeAll("print"),
117 .f_div => try writer.writeAll("/."),
118 .f_mul => try writer.writeAll("*."),
119 .f_add => try writer.writeAll("+."),
120 .f_sub => try writer.writeAll("-."),
121 .error_op => try writer.writeAll("error"),
122 }
123 }
124
125 pub fn primArity(self: Head) ?usize {
126 return switch (self) {
127 .y_combinator => 1,
128 .flip => 1,
129 .factor => 1,
130 .native_eq => 2,
131 .get_args => 1,
132 .get_constructor => 1,
133 .pbool => 1,
134 .get_config => 0,
135 .mk_int => 2,
136 .mk_int_weighted => 2,
137 .int_dist_eq => 2,
138 .print_op => 1,
139 .f_div => 2,
140 .f_mul => 2,
141 .f_add => 2,
142 .f_sub => 2,
143 .error_op => 1,
144 else => null,
145 };
146 }
147 };
148
149 pub const PExpr = struct {
150 head: Head,
151 args: []const *PExpr,
152
153 const Self = @This();
154
155 pub fn init(allocator: Allocator, head: Head) !*Self {
156 return initWithArgs(allocator, head, &[_]*Self{});
157 }
158
159 pub fn initWithArgs(allocator: Allocator, head: Head, args: []const *Self) !*Self {
160 const self = try allocator.create(Self);
161 const args_copy = try allocator.alloc(*Self, args.len);
162 @memcpy(args_copy, args);
163 self.* = Self{
164 .head = head,
165 .args = args_copy,
166 };
167 return self;
168 }
169
170 pub fn deinit(self: *Self, allocator: Allocator) void {
171 for (self.args) |arg| {
172 arg.deinit(allocator);
173 }
174 switch (self.head) {
175 .case_of => |c| {
176 for (c.branches) |b| {
177 allocator.free(b.args);
178 }
179 allocator.free(c.branches);
180 },
181 else => {},
182 }
183 allocator.free(self.args);
184 allocator.destroy(self);
185 }
186
187 pub fn format(
188 self: *const Self,
189 comptime fmt: []const u8,
190 options: std.fmt.Options,
191 writer: anytype,
192 ) !void {
193 switch (self.head) {
194 .var_ref => |v| {
195 try writer.print("${s}", .{v.name});
196 return;
197 },
198 .defined => |d| {
199 try writer.print("{s}", .{d.name});
200 return;
201 },
202 .const_native => |c| {
203 try c.format(fmt, options, writer);
204 return;
205 },
206 .abs => |a| {
207 try writer.print("(λ{s}", .{a.var_name});
208 var body = self.args[0];
209 while (body.head == .abs) {
210 const inner_abs = body.head.abs;
211 try writer.print(" {s}", .{inner_abs.var_name});
212 body = body.args[0];
213 }
214 try writer.writeAll(" -> ");
215 try body.format(fmt, options, writer);
216 try writer.writeAll(")");
217 return;
218 },
219 .app => {
220 if (self.args[0].head == .abs) {
221 try self.formatLetApplication(fmt, options, writer);
222 return;
223 }
224
225 try writer.writeAll("(");
226 try self.formatApplicationChain(fmt, options, writer);
227 try writer.writeAll(")");
228 return;
229 },
230 .case_of => |c| {
231 if (c.branches.len == 2 and
232 std.mem.eql(u8, c.branches[0].constructor, "True") and
233 std.mem.eql(u8, c.branches[1].constructor, "False"))
234 {
235 try writer.writeAll("(if ");
236 try self.args[0].format(fmt, options, writer);
237 try writer.writeAll(" ");
238 try self.args[1].format(fmt, options, writer);
239 try writer.writeAll(" ");
240 try self.args[2].format(fmt, options, writer);
241 try writer.writeAll(")");
242 return;
243 }
244
245 try writer.writeAll("(case ");
246 try self.args[0].format(fmt, options, writer);
247 try writer.writeAll(" of ");
248 for (c.branches, 0..) |branch, idx| {
249 try branch.format(fmt, options, writer);
250 try writer.writeAll(" => ");
251 try self.args[idx + 1].format(fmt, options, writer);
252 if (idx < c.branches.len - 1) try writer.writeAll(" | ");
253 }
254 try writer.writeAll(")");
255 return;
256 },
257 .construct => |c| {
258 if (self.maybeConst()) |n| {
259 try writer.print("{d}", .{n});
260 return;
261 }
262
263 try writer.print("({s}", .{c.constructor});
264 for (self.args) |arg| {
265 try writer.writeAll(" ");
266 try arg.format(fmt, options, writer);
267 }
268 try writer.writeAll(")");
269 return;
270 },
271 else => {},
272 }
273
274 try writer.writeAll("(");
275 try self.head.format(fmt, options, writer);
276 for (self.args) |arg| {
277 try writer.writeAll(" ");
278 try arg.format(fmt, options, writer);
279 }
280 try writer.writeAll(")");
281 }
282
283 fn formatLetApplication(
284 self: *const Self,
285 comptime fmt: []const u8,
286 options: std.fmt.Options,
287 writer: anytype,
288 ) anyerror!void {
289 try writer.writeAll("(let [");
290 var current = self;
291 var first = true;
292 while (current.head == .app and current.args[0].head == .abs) {
293 if (!first) try writer.writeAll(" ");
294 const abs_head = current.args[0].head.abs;
295 try writer.print("{s} ", .{abs_head.var_name});
296 try current.args[1].format(fmt, options, writer);
297 first = false;
298 current = current.args[0].args[0];
299 }
300 try writer.writeAll("] ");
301 try current.format(fmt, options, writer);
302 try writer.writeAll(")");
303 }
304
305 fn formatApplicationChain(
306 self: *const Self,
307 comptime fmt: []const u8,
308 options: std.fmt.Options,
309 writer: anytype,
310 ) anyerror!void {
311 if (self.head != .app) {
312 try self.format(fmt, options, writer);
313 return;
314 }
315
316 try self.args[0].formatApplicationChain(fmt, options, writer);
317 try writer.writeAll(" ");
318 try self.args[1].format(fmt, options, writer);
319 }
320
321 pub fn maybeConst(self: *const Self) ?i64 {
322 switch (self.head) {
323 .construct => |c| {
324 if (std.mem.eql(u8, c.constructor, "O")) {
325 return 0;
326 } else if (std.mem.eql(u8, c.constructor, "S")) {
327 if (self.args.len == 1) {
328 if (self.args[0].maybeConst()) |inner| {
329 return inner + 1;
330 }
331 }
332 }
333 return null;
334 },
335 else => return null,
336 }
337 }
338 };
339
340 pub const TypeRegistry = struct {
341 allocator: Allocator,
342 type_of_constructor: std.StringHashMap(Symbol),
343 constructors_of_type: std.StringHashMap([]const Symbol),
344 args_of_constructor: std.StringHashMap([]const Symbol),
345
346 const Self = @This();
347
348 pub fn init(allocator: Allocator) Self {
349 return Self{
350 .allocator = allocator,
351 .type_of_constructor = std.StringHashMap(Symbol).init(allocator),
352 .constructors_of_type = std.StringHashMap([]const Symbol).init(allocator),
353 .args_of_constructor = std.StringHashMap([]const Symbol).init(allocator),
354 };
355 }
356
357 pub fn deinit(self: *Self) void {
358 var iter = self.constructors_of_type.iterator();
359 while (iter.next()) |entry| {
360 self.allocator.free(entry.value_ptr.*);
361 }
362 var iter2 = self.args_of_constructor.iterator();
363 while (iter2.next()) |entry| {
364 self.allocator.free(entry.value_ptr.*);
365 }
366 self.type_of_constructor.deinit();
367 self.constructors_of_type.deinit();
368 self.args_of_constructor.deinit();
369 }
370
371 pub fn defineType(
372 self: *Self,
373 type_name: Symbol,
374 constructors: []const struct { name: Symbol, args: []const Symbol },
375 ) !void {
376 var constructor_names = try self.allocator.alloc(Symbol, constructors.len);
377 for (constructors, 0..) |ctor, i| {
378 try self.type_of_constructor.put(ctor.name, type_name);
379 const args_copy = try self.allocator.alloc(Symbol, ctor.args.len);
380 @memcpy(args_copy, ctor.args);
381 try self.args_of_constructor.put(ctor.name, args_copy);
382 constructor_names[i] = ctor.name;
383 }
384 try self.constructors_of_type.put(type_name, constructor_names);
385 }
386
387 pub fn hasConstructor(self: *const Self, name: Symbol) bool {
388 return self.args_of_constructor.contains(name);
389 }
390
391 pub fn constructorArity(self: *const Self, name: Symbol) ?usize {
392 if (self.args_of_constructor.get(name)) |args| {
393 return args.len;
394 }
395 return null;
396 }
397
398 pub fn initWithDefaults(allocator: Allocator) !Self {
399 var self = Self.init(allocator);
400
401 try self.defineType("nat", &.{
402 .{ .name = "O", .args = &.{} },
403 .{ .name = "S", .args = &.{"nat"} },
404 });
405
406 try self.defineType("list", &.{
407 .{ .name = "Nil", .args = &.{} },
408 .{ .name = "Cons", .args = &.{ "nat", "list" } },
409 });
410
411 try self.defineType("snoclist", &.{
412 .{ .name = "SNil", .args = &.{} },
413 .{ .name = "Snoc", .args = &.{ "snoclist", "nat" } },
414 });
415
416 try self.defineType("bool", &.{
417 .{ .name = "True", .args = &.{} },
418 .{ .name = "False", .args = &.{} },
419 });
420
421 try self.defineType("unit", &.{
422 .{ .name = "Unit", .args = &.{} },
423 });
424
425 return self;
426 }
427 };
428
429 pub const Definition = struct {
430 name: Symbol,
431 expr: *PExpr,
432 is_stdlib: bool = false,
433 doc: ?[]const u8 = null,
434 };
435
436 pub const Definitions = struct {
437 allocator: Allocator,
438 defs: std.StringHashMap(Definition),
439
440 const Self = @This();
441
442 pub const DefineOptions = struct {
443 is_stdlib: bool = false,
444 doc: ?[]const u8 = null,
445 };
446
447 pub fn init(allocator: Allocator) Self {
448 return Self{
449 .allocator = allocator,
450 .defs = std.StringHashMap(Definition).init(allocator),
451 };
452 }
453
454 pub fn deinit(self: *Self) void {
455 var iter = self.defs.iterator();
456 while (iter.next()) |entry| {
457 entry.value_ptr.expr.deinit(self.allocator);
458 if (entry.value_ptr.doc) |doc| {
459 self.allocator.free(doc);
460 }
461 }
462 self.defs.deinit();
463 }
464
465 pub fn define(self: *Self, name: Symbol, expr: *PExpr) !void {
466 try self.defineWithOptions(name, expr, .{});
467 }
468
469 pub fn defineStdlib(self: *Self, name: Symbol, expr: *PExpr) !void {
470 try self.defineWithOptions(name, expr, .{ .is_stdlib = true });
471 }
472
473 pub fn defineWithDoc(self: *Self, name: Symbol, expr: *PExpr, doc: ?[]const u8) !void {
474 try self.defineWithOptions(name, expr, .{ .doc = doc });
475 }
476
477 pub fn defineStdlibWithDoc(self: *Self, name: Symbol, expr: *PExpr, doc: ?[]const u8) !void {
478 try self.defineWithOptions(name, expr, .{ .is_stdlib = true, .doc = doc });
479 }
480
481 fn defineWithOptions(self: *Self, name: Symbol, expr: *PExpr, options: DefineOptions) !void {
482 if (self.defs.get(name)) |existing| {
483 existing.expr.deinit(self.allocator);
484 if (existing.doc) |doc| {
485 self.allocator.free(doc);
486 }
487 }
488 const doc_copy: ?[]const u8 = if (options.doc) |d| try self.allocator.dupe(u8, d) else null;
489 try self.defs.put(name, Definition{
490 .name = name,
491 .expr = expr,
492 .is_stdlib = options.is_stdlib,
493 .doc = doc_copy,
494 });
495 }
496
497 pub fn lookup(self: *const Self, name: Symbol) ?*PExpr {
498 if (self.defs.get(name)) |def| {
499 return def.expr;
500 }
501 return null;
502 }
503
504 pub fn lookupDefinition(self: *const Self, name: Symbol) ?Definition {
505 return self.defs.get(name);
506 }
507
508 pub fn isDefined(self: *const Self, name: Symbol) bool {
509 return self.defs.contains(name);
510 }
511
512 pub fn remove(self: *Self, name: Symbol) ?*PExpr {
513 if (self.defs.fetchRemove(name)) |kv| {
514 return kv.value.expr;
515 }
516 return null;
517 }
518
519 pub fn clearUserDefinitions(self: *Self) void {
520 var to_remove: std.ArrayList(Symbol) = .empty;
521 defer to_remove.deinit(self.allocator);
522
523 var iter = self.defs.iterator();
524 while (iter.next()) |entry| {
525 if (!entry.value_ptr.is_stdlib) {
526 to_remove.append(self.allocator, entry.key_ptr.*) catch continue;
527 }
528 }
529
530 for (to_remove.items) |name| {
531 if (self.defs.fetchRemove(name)) |kv| {
532 kv.value.expr.deinit(self.allocator);
533 }
534 }
535 }
536 };
537
538 pub const Token = []const u8;
539
540 fn isTokenDelimiter(c: u8) bool {
541 return switch (c) {
542 '(', ')', '{', '}', '[', ']', ',', '~', '`' => true,
543 else => false,
544 };
545 }
546
547 fn isTokenWhitespace(c: u8) bool {
548 return switch (c) {
549 ' ', '\t', '\r', '\n' => true,
550 else => false,
551 };
552 }
553
554 fn skipIgnored(source: []const u8, index: *usize) void {
555 while (index.* < source.len) {
556 const i = index.*;
557 const c = source[i];
558
559 if (c == ';' and i + 1 < source.len and source[i + 1] == ';') {
560 index.* += 2;
561 while (index.* < source.len and source[index.*] != '\n') {
562 index.* += 1;
563 }
564 continue;
565 }
566
567 if (isTokenWhitespace(c)) {
568 index.* += 1;
569 continue;
570 }
571
572 break;
573 }
574 }
575
576 fn isTwoByteTokenAt(source: []const u8, i: usize) bool {
577 if (i + 1 >= source.len) return false;
578 const c = source[i];
579 const next = source[i + 1];
580 return (c == '-' and next == '>') or
581 (c == '=' and next == '>') or
582 (c == 0xCE and next == 0xBB);
583 }
584
585 fn nextTokenSpan(source: []const u8, index: *usize) ?struct { start: usize, end: usize } {
586 skipIgnored(source, index);
587 if (index.* >= source.len) return null;
588
589 const start = index.*;
590 const c = source[start];
591
592 if (isTwoByteTokenAt(source, start)) {
593 index.* += 2;
594 return .{ .start = start, .end = index.* };
595 }
596
597 if (isTokenDelimiter(c) or c == '|') {
598 index.* += 1;
599 return .{ .start = start, .end = index.* };
600 }
601
602 while (index.* < source.len) {
603 const i = index.*;
604 const current = source[i];
605 if (isTokenWhitespace(current)) break;
606 if (current == ';' and i + 1 < source.len and source[i + 1] == ';') break;
607 if (isTwoByteTokenAt(source, i)) break;
608 if (isTokenDelimiter(current) or current == '|') break;
609 index.* += 1;
610 }
611
612 return .{ .start = start, .end = index.* };
613 }
614
615 pub fn tokenize(allocator: Allocator, source: []const u8) ![]Token {
616 var tokens: std.ArrayList(Token) = .empty;
617 errdefer tokens.deinit(allocator);
618
619 var index: usize = 0;
620 while (nextTokenSpan(source, &index)) |span| {
621 try tokens.append(allocator, source[span.start..span.end]);
622 }
623
624 return try tokens.toOwnedSlice(allocator);
625 }
626
627 pub fn freeTokens(allocator: Allocator, tokens: []Token) void {
628 allocator.free(tokens);
629 }
630
631 pub const ParseError = error{
632 UnexpectedEndOfInput,
633 ExpectedClosingParen,
634 ExpectedClosingBracket,
635 ExpectedArrow,
636 InvalidIdentifier,
637 InvalidExpression,
638 UnknownToken,
639 WrongArgumentCount,
640 DuplicateConstructor,
641 ExpectedOf,
642 OutOfMemory,
643 };
644
645 pub const Parser = struct {
646 allocator: Allocator,
647 tokens: []const Token,
648 pos: usize,
649 env: std.ArrayList(Symbol),
650 types: *const TypeRegistry,
651 defs: *const Definitions,
652
653 const Self = @This();
654
655 pub fn init(
656 allocator: Allocator,
657 tokens: []const Token,
658 types: *const TypeRegistry,
659 defs: *const Definitions,
660 ) Self {
661 return Self{
662 .allocator = allocator,
663 .tokens = tokens,
664 .pos = 0,
665 .env = .empty,
666 .types = types,
667 .defs = defs,
668 };
669 }
670
671 pub fn deinit(self: *Self) void {
672 self.env.deinit(self.allocator);
673 }
674
675 pub fn peek(self: *const Self) ?Token {
676 if (self.pos < self.tokens.len) {
677 return self.tokens[self.pos];
678 }
679 return null;
680 }
681
682 fn advance(self: *Self) ?Token {
683 if (self.pos < self.tokens.len) {
684 const token = self.tokens[self.pos];
685 self.pos += 1;
686 return token;
687 }
688 return null;
689 }
690
691 fn expect(self: *Self, expected: []const u8) !void {
692 const token = self.advance() orelse return ParseError.UnexpectedEndOfInput;
693 if (!std.mem.eql(u8, token, expected)) {
694 return ParseError.ExpectedClosingParen;
695 }
696 }
697
698 pub fn pushEnv(self: *Self, name: Symbol) !void {
699 try self.env.append(self.allocator, name);
700 }
701
702 fn dupeStr(self: *Self, s: []const u8) ![]const u8 {
703 return try self.allocator.dupe(u8, s);
704 }
705
706 pub fn popEnv(self: *Self) void {
707 _ = self.env.pop();
708 }
709
710 fn inEnv(self: *const Self, name: Symbol) bool {
711 for (self.env.items) |item| {
712 if (std.mem.eql(u8, item, name)) return true;
713 }
714 return false;
715 }
716
717 pub fn parseExpr(self: *Self) ParseError!*PExpr {
718 const token = self.advance() orelse return ParseError.UnexpectedEndOfInput;
719
720 if (std.mem.eql(u8, token, "(")) {
721 return self.parseCompound();
722 }
723
724 if (std.mem.eql(u8, token, "[")) {
725 return self.parseList();
726 }
727
728 if (token.len > 1 and token[0] == '\'') {
729 return PExpr.init(self.allocator, .{
730 .const_native = .{ .symbol = try self.dupeStr(token[1..]) },
731 });
732 }
733
734 if (token.len > 1 and token[0] == '@') {
735 const val = std.fmt.parseInt(i64, token[1..], 10) catch return ParseError.InvalidIdentifier;
736 return PExpr.init(self.allocator, .{
737 .const_native = .{ .int = val },
738 });
739 }
740
741 if (isInteger(token)) {
742 const val = std.fmt.parseInt(i64, token, 10) catch return ParseError.InvalidIdentifier;
743 return self.constToExpr(val);
744 }
745
746 if (isFloat(token)) {
747 const val = std.fmt.parseFloat(f64, token) catch return ParseError.InvalidIdentifier;
748 return PExpr.init(self.allocator, .{
749 .const_native = .{ .float = val },
750 });
751 }
752
753 if (std.mem.eql(u8, token, "true")) {
754 return PExpr.init(self.allocator, .{ .construct = .{ .constructor = "True" } });
755 }
756 if (std.mem.eql(u8, token, "false")) {
757 return PExpr.init(self.allocator, .{ .construct = .{ .constructor = "False" } });
758 }
759
760 if (std.mem.eql(u8, token, "nothing")) {
761 return PExpr.init(self.allocator, .{ .construct = .{ .constructor = "Unit" } });
762 }
763
764 if (token.len > 1 and token[0] == '$') {
765 return PExpr.init(self.allocator, .{ .var_ref = .{ .name = try self.dupeStr(token[1..]) } });
766 }
767
768 if (self.inEnv(token)) {
769 return PExpr.init(self.allocator, .{ .var_ref = .{ .name = try self.dupeStr(token) } });
770 }
771
772 if (self.defs.isDefined(token)) {
773 return PExpr.init(self.allocator, .{ .defined = .{ .name = try self.dupeStr(token) } });
774 }
775
776 if (self.types.hasConstructor(token)) {
777 if (self.types.constructorArity(token) == 0) {
778 return PExpr.init(self.allocator, .{ .construct = .{ .constructor = try self.dupeStr(token) } });
779 }
780 }
781
782 return ParseError.UnknownToken;
783 }
784
785 fn parseCompound(self: *Self) ParseError!*PExpr {
786 const head_token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
787
788 if (isLambdaKeyword(head_token)) {
789 _ = self.advance();
790 return self.parseLambda();
791 }
792
793 if (std.mem.eql(u8, head_token, "if")) {
794 _ = self.advance();
795 return self.parseIf();
796 }
797
798 if (std.mem.eql(u8, head_token, "Y")) {
799 _ = self.advance();
800 return self.parseY();
801 }
802
803 if (std.mem.eql(u8, head_token, "case") or std.mem.eql(u8, head_token, "match")) {
804 _ = self.advance();
805 return self.parseCase();
806 }
807
808 if (std.mem.eql(u8, head_token, "let")) {
809 _ = self.advance();
810 return self.parseLet();
811 }
812
813 if (self.types.hasConstructor(head_token)) {
814 _ = self.advance();
815 return self.parseConstructor(head_token);
816 }
817
818 if (lookupPrim(head_token)) |head| {
819 _ = self.advance();
820 return self.parsePrimitive(head);
821 }
822
823 if (std.mem.eql(u8, head_token, "discrete")) {
824 _ = self.advance();
825 return self.parseDiscrete();
826 }
827
828 if (std.mem.eql(u8, head_token, "uniform")) {
829 _ = self.advance();
830 return self.parseUniform();
831 }
832
833 return self.parseApplication();
834 }
835
836 fn parseLambda(self: *Self) ParseError!*PExpr {
837 var arg_names: std.ArrayList(Symbol) = .empty;
838 defer arg_names.deinit(self.allocator);
839
840 const first = self.peek() orelse return ParseError.UnexpectedEndOfInput;
841 if (std.mem.eql(u8, first, "->")) {
842 _ = self.advance();
843 try self.pushEnv("_");
844 const body = try self.parseExpr();
845 self.popEnv();
846 try self.expect(")");
847 return PExpr.initWithArgs(self.allocator, .{ .abs = .{ .var_name = "_" } }, &.{body});
848 }
849
850 while (true) {
851 const token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
852
853 if (std.mem.eql(u8, token, "->")) {
854 _ = self.advance();
855 break;
856 }
857
858 if (std.mem.eql(u8, token, ",")) {
859 _ = self.advance();
860 continue;
861 }
862
863 if (!isIdentifier(token)) return ParseError.InvalidIdentifier;
864 _ = self.advance();
865
866 const duped_name = try self.dupeStr(token);
867 try arg_names.append(self.allocator, duped_name);
868 try self.pushEnv(token);
869 }
870
871 const body = try self.parseExpr();
872 try self.expect(")");
873
874 for (arg_names.items) |_| {
875 self.popEnv();
876 }
877
878 var result = body;
879 var i = arg_names.items.len;
880 while (i > 0) {
881 i -= 1;
882 result = try PExpr.initWithArgs(self.allocator, .{ .abs = .{ .var_name = arg_names.items[i] } }, &.{result});
883 }
884
885 return result;
886 }
887
888 fn parseIf(self: *Self) ParseError!*PExpr {
889 const cond = try self.parseExpr();
890 const then_expr = try self.parseExpr();
891 const else_expr = try self.parseExpr();
892 try self.expect(")");
893
894 const branches = try self.allocator.alloc(CaseOfGuard, 2);
895 branches[0] = CaseOfGuard{ .constructor = "True", .args = &.{} };
896 branches[1] = CaseOfGuard{ .constructor = "False", .args = &.{} };
897
898 return PExpr.initWithArgs(
899 self.allocator,
900 .{ .case_of = .{ .branches = branches } },
901 &.{ cond, then_expr, else_expr },
902 );
903 }
904
905 fn parseY(self: *Self) ParseError!*PExpr {
906 const f = try self.parseExpr();
907
908 const next = self.peek();
909 if (next) |n| {
910 if (!std.mem.eql(u8, n, ")")) {
911 const x = try self.parseExpr();
912 try self.expect(")");
913
914 const y_expr = try PExpr.initWithArgs(self.allocator, .y_combinator, &.{f});
915 return PExpr.initWithArgs(self.allocator, .app, &.{ y_expr, x });
916 }
917 }
918
919 try self.expect(")");
920 return PExpr.initWithArgs(self.allocator, .y_combinator, &.{f});
921 }
922
923 fn parseCase(self: *Self) ParseError!*PExpr {
924 const scrutinee = try self.parseExpr();
925
926 const maybe_of = self.peek();
927 if (maybe_of) |token| {
928 if (std.mem.eql(u8, token, "of")) {
929 _ = self.advance();
930 }
931 }
932
933 var guards: std.ArrayList(CaseOfGuard) = .empty;
934 var branches: std.ArrayList(*PExpr) = .empty;
935 defer guards.deinit(self.allocator);
936 defer branches.deinit(self.allocator);
937
938 while (true) {
939 const token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
940 if (std.mem.eql(u8, token, ")")) {
941 _ = self.advance();
942 break;
943 }
944
945 const constructor_tok = self.advance() orelse return ParseError.UnexpectedEndOfInput;
946 const constructor = try self.dupeStr(constructor_tok);
947
948 var args: std.ArrayList(Symbol) = .empty;
949 while (true) {
950 const arg_token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
951 if (std.mem.eql(u8, arg_token, "=>")) {
952 _ = self.advance();
953 break;
954 }
955 _ = self.advance();
956 const duped_arg = try self.dupeStr(arg_token);
957 try args.append(self.allocator, duped_arg);
958 try self.pushEnv(arg_token);
959 }
960
961 for (guards.items) |g| {
962 if (std.mem.eql(u8, g.constructor, constructor)) {
963 return ParseError.DuplicateConstructor;
964 }
965 }
966
967 const args_slice = try args.toOwnedSlice(self.allocator);
968 try guards.append(self.allocator, CaseOfGuard{ .constructor = constructor, .args = args_slice });
969
970 const body = try self.parseExpr();
971 try branches.append(self.allocator, body);
972
973 for (args_slice) |_| {
974 self.popEnv();
975 }
976
977 const sep = self.peek();
978 if (sep) |s| {
979 if (std.mem.eql(u8, s, "|")) {
980 _ = self.advance();
981 }
982 }
983 }
984
985 const guards_slice = try guards.toOwnedSlice(self.allocator);
986
987 if (guards_slice.len == 0) {
988 return ParseError.InvalidExpression;
989 }
990
991 var all_args: std.ArrayList(*PExpr) = .empty;
992 defer all_args.deinit(self.allocator);
993 try all_args.append(self.allocator, scrutinee);
994 try all_args.appendSlice(self.allocator, branches.items);
995
996 return PExpr.initWithArgs(
997 self.allocator,
998 .{ .case_of = .{ .branches = guards_slice } },
999 try all_args.toOwnedSlice(self.allocator),
1000 );
1001 }
1002
1003 fn parseLet(self: *Self) ParseError!*PExpr {
1004 const open = self.advance() orelse return ParseError.UnexpectedEndOfInput;
1005 const close_token: []const u8 = if (std.mem.eql(u8, open, "[")) "]" else ")";
1006
1007 const Binding = struct { name: Symbol, val: *PExpr };
1008 var bindings: std.ArrayList(Binding) = .empty;
1009 defer bindings.deinit(self.allocator);
1010
1011 while (true) {
1012 const token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
1013
1014 if (std.mem.eql(u8, token, close_token)) {
1015 _ = self.advance();
1016 break;
1017 }
1018
1019 if (std.mem.eql(u8, token, "(")) {
1020 _ = self.advance();
1021 const var_name_tok = self.advance() orelse return ParseError.UnexpectedEndOfInput;
1022 const var_name = try self.dupeStr(var_name_tok);
1023 const val = try self.parseExpr();
1024 try self.expect(")");
1025 try bindings.append(self.allocator, .{ .name = var_name, .val = val });
1026 try self.pushEnv(var_name);
1027 } else {
1028 const var_name_tok = self.advance() orelse return ParseError.UnexpectedEndOfInput;
1029 const var_name = try self.dupeStr(var_name_tok);
1030 const val = try self.parseExpr();
1031 try bindings.append(self.allocator, .{ .name = var_name, .val = val });
1032 try self.pushEnv(var_name);
1033 }
1034 }
1035
1036 const body = try self.parseExpr();
1037 try self.expect(")");
1038
1039 for (bindings.items) |_| {
1040 self.popEnv();
1041 }
1042
1043 var result = body;
1044 var i = bindings.items.len;
1045 while (i > 0) {
1046 i -= 1;
1047 const binding = bindings.items[i];
1048 const abs_expr = try PExpr.initWithArgs(self.allocator, .{ .abs = .{ .var_name = binding.name } }, &.{result});
1049 result = try PExpr.initWithArgs(self.allocator, .app, &.{ abs_expr, binding.val });
1050 }
1051
1052 return result;
1053 }
1054
1055 fn parseConstructor(self: *Self, constructor_tok: Symbol) ParseError!*PExpr {
1056 const constructor = try self.dupeStr(constructor_tok);
1057
1058 var args: std.ArrayList(*PExpr) = .empty;
1059 defer args.deinit(self.allocator);
1060
1061 while (true) {
1062 const token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
1063 if (std.mem.eql(u8, token, ")")) {
1064 _ = self.advance();
1065 break;
1066 }
1067 const arg = try self.parseExpr();
1068 try args.append(self.allocator, arg);
1069 }
1070
1071 return PExpr.initWithArgs(
1072 self.allocator,
1073 .{ .construct = .{ .constructor = constructor } },
1074 try args.toOwnedSlice(self.allocator),
1075 );
1076 }
1077
1078 fn parsePrimitive(self: *Self, head: Head) ParseError!*PExpr {
1079 const arity = head.primArity() orelse 0;
1080
1081 var args: std.ArrayList(*PExpr) = .empty;
1082 defer args.deinit(self.allocator);
1083
1084 for (0..arity) |_| {
1085 const arg = try self.parseExpr();
1086 try args.append(self.allocator, arg);
1087 }
1088
1089 try self.expect(")");
1090
1091 return PExpr.initWithArgs(self.allocator, head, try args.toOwnedSlice(self.allocator));
1092 }
1093
1094 fn parseApplication(self: *Self) ParseError!*PExpr {
1095 const func = try self.parseExpr();
1096
1097 var args: std.ArrayList(*PExpr) = .empty;
1098 defer args.deinit(self.allocator);
1099
1100 while (true) {
1101 const token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
1102 if (std.mem.eql(u8, token, ")")) {
1103 _ = self.advance();
1104 break;
1105 }
1106 const arg = try self.parseExpr();
1107 try args.append(self.allocator, arg);
1108 }
1109
1110 if (args.items.len == 0) {
1111 const unit = try PExpr.init(self.allocator, .{ .construct = .{ .constructor = "Unit" } });
1112 try args.append(self.allocator, unit);
1113 }
1114
1115 var result = func;
1116 for (args.items) |arg| {
1117 result = try PExpr.initWithArgs(self.allocator, .app, &.{ result, arg });
1118 }
1119
1120 return result;
1121 }
1122
1123 fn parseList(self: *Self) ParseError!*PExpr {
1124 var items: std.ArrayList(*PExpr) = .empty;
1125 defer items.deinit(self.allocator);
1126
1127 while (true) {
1128 const token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
1129 if (std.mem.eql(u8, token, "]")) {
1130 _ = self.advance();
1131 break;
1132 }
1133 if (std.mem.eql(u8, token, ",")) {
1134 _ = self.advance();
1135 continue;
1136 }
1137 const item = try self.parseExpr();
1138 try items.append(self.allocator, item);
1139 }
1140
1141 var result = try PExpr.init(self.allocator, .{ .construct = .{ .constructor = "Nil" } });
1142 var i = items.items.len;
1143 while (i > 0) {
1144 i -= 1;
1145 result = try PExpr.initWithArgs(
1146 self.allocator,
1147 .{ .construct = .{ .constructor = "Cons" } },
1148 &.{ items.items[i], result },
1149 );
1150 }
1151
1152 return result;
1153 }
1154
1155 fn parseDiscrete(self: *Self) ParseError!*PExpr {
1156 var options: std.ArrayList(*PExpr) = .empty;
1157 var probs: std.ArrayList(f64) = .empty;
1158 defer options.deinit(self.allocator);
1159 defer probs.deinit(self.allocator);
1160
1161 while (true) {
1162 const token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
1163 if (std.mem.eql(u8, token, ")")) {
1164 _ = self.advance();
1165 break;
1166 }
1167
1168 try self.expect("(");
1169 const expr = try self.parseExpr();
1170 const prob_token = self.advance() orelse return ParseError.UnexpectedEndOfInput;
1171 const prob = std.fmt.parseFloat(f64, prob_token) catch return ParseError.InvalidIdentifier;
1172 try self.expect(")");
1173
1174 try options.append(self.allocator, expr);
1175 try probs.append(self.allocator, prob);
1176 }
1177
1178 if (options.items.len == 0) {
1179 return ParseError.InvalidExpression;
1180 }
1181
1182 var filtered_opts: std.ArrayList(*PExpr) = .empty;
1183 var filtered_probs: std.ArrayList(f64) = .empty;
1184 defer filtered_opts.deinit(self.allocator);
1185 defer filtered_probs.deinit(self.allocator);
1186
1187 var total: f64 = 0.0;
1188 for (options.items, probs.items) |opt, p| {
1189 if (p > 0.0) {
1190 try filtered_opts.append(self.allocator, opt);
1191 try filtered_probs.append(self.allocator, p);
1192 total += p;
1193 }
1194 }
1195
1196 if (filtered_opts.items.len == 0) {
1197 if (!builtin.is_test) {
1198 log.warn("discrete: all probabilities are zero or negative", .{});
1199 }
1200 return ParseError.InvalidExpression;
1201 }
1202
1203 if (@abs(total - 1.0) > 1e-5) {
1204 if (!builtin.is_test) {
1205 log.warn("discrete: probabilities sum to {d:.6} (must be 1.0)", .{total});
1206 }
1207 return ParseError.InvalidExpression;
1208 }
1209
1210 return self.buildDiscrete(filtered_opts.items, filtered_probs.items);
1211 }
1212
1213 fn parseUniform(self: *Self) ParseError!*PExpr {
1214 var options: std.ArrayList(*PExpr) = .empty;
1215 defer options.deinit(self.allocator);
1216
1217 while (true) {
1218 const token = self.peek() orelse return ParseError.UnexpectedEndOfInput;
1219 if (std.mem.eql(u8, token, ")")) {
1220 _ = self.advance();
1221 break;
1222 }
1223 const expr = try self.parseExpr();
1224 try options.append(self.allocator, expr);
1225 }
1226
1227 if (options.items.len == 0) {
1228 if (!builtin.is_test) {
1229 log.warn("uniform: requires at least one option", .{});
1230 }
1231 return ParseError.InvalidExpression;
1232 }
1233
1234 const n = options.items.len;
1235 const probs = try self.allocator.alloc(f64, n);
1236 defer self.allocator.free(probs);
1237 for (probs) |*p| {
1238 p.* = 1.0 / @as(f64, @floatFromInt(n));
1239 }
1240
1241 return self.buildDiscrete(options.items, probs);
1242 }
1243
1244 fn buildDiscrete(self: *Self, options: []*PExpr, probs: []f64) ParseError!*PExpr {
1245 const n = options.len;
1246
1247 if (n == 0) return ParseError.InvalidExpression;
1248
1249 if (n == 1) return options[0];
1250
1251 if (n == 2) {
1252 const flip_arg = try PExpr.init(self.allocator, .{ .const_native = .{ .float = probs[1] } });
1253 const flip_expr = try PExpr.initWithArgs(self.allocator, .flip, &.{flip_arg});
1254
1255 const branches = try self.allocator.alloc(CaseOfGuard, 2);
1256 branches[0] = CaseOfGuard{ .constructor = "True", .args = &.{} };
1257 branches[1] = CaseOfGuard{ .constructor = "False", .args = &.{} };
1258
1259 return PExpr.initWithArgs(
1260 self.allocator,
1261 .{ .case_of = .{ .branches = branches } },
1262 &.{ flip_expr, options[1], options[0] },
1263 );
1264 }
1265
1266 const num_bits = std.math.log2_int_ceil(usize, n);
1267 const padded_len = @as(usize, 1) << @intCast(num_bits);
1268
1269 const padded_probs = try self.allocator.alloc(f64, padded_len);
1270 defer self.allocator.free(padded_probs);
1271
1272 for (padded_probs, 0..) |*p, i| {
1273 p.* = if (i < n) probs[i] else 0.0;
1274 }
1275
1276 return self.buildDiscreteRecursive(options, padded_probs, num_bits, 0, padded_len);
1277 }
1278
1279 fn buildDiscreteRecursive(
1280 self: *Self,
1281 options: []*PExpr,
1282 probs: []f64,
1283 bit_idx: usize,
1284 start: usize,
1285 end: usize,
1286 ) ParseError!*PExpr {
1287 if (bit_idx == 0) {
1288 return if (start < options.len) options[start] else options[0];
1289 }
1290
1291 const mid = (start + end) / 2;
1292
1293 var denom: f64 = 0;
1294 for (probs[start..end]) |p| denom += p;
1295
1296 if (denom == 0) {
1297 return self.buildDiscreteRecursive(options, probs, bit_idx - 1, start, mid);
1298 }
1299
1300 var right_sum: f64 = 0;
1301 for (probs[mid..end]) |p| right_sum += p;
1302 const p = right_sum / denom;
1303
1304 const left = try self.buildDiscreteRecursive(options, probs, bit_idx - 1, mid, end);
1305 const right = try self.buildDiscreteRecursive(options, probs, bit_idx - 1, start, mid);
1306
1307 const flip_arg = try PExpr.init(self.allocator, .{ .const_native = .{ .float = p } });
1308 const flip_expr = try PExpr.initWithArgs(self.allocator, .flip, &.{flip_arg});
1309
1310 const branches = try self.allocator.alloc(CaseOfGuard, 2);
1311 branches[0] = CaseOfGuard{ .constructor = "True", .args = &.{} };
1312 branches[1] = CaseOfGuard{ .constructor = "False", .args = &.{} };
1313
1314 return PExpr.initWithArgs(
1315 self.allocator,
1316 .{ .case_of = .{ .branches = branches } },
1317 &.{ flip_expr, left, right },
1318 );
1319 }
1320
1321 fn constToExpr(self: *Self, val: i64) ParseError!*PExpr {
1322 var result = try PExpr.init(self.allocator, .{ .construct = .{ .constructor = "O" } });
1323 var i: i64 = 0;
1324 while (i < val) : (i += 1) {
1325 result = try PExpr.initWithArgs(
1326 self.allocator,
1327 .{ .construct = .{ .constructor = "S" } },
1328 &.{result},
1329 );
1330 }
1331 return result;
1332 }
1333 };
1334
1335 fn lookupPrim(name: []const u8) ?Head {
1336 const prims = .{
1337 .{ "Y", Head.y_combinator },
1338 .{ "flip", Head.flip },
1339 .{ "factor", Head.factor },
1340 .{ "native_eq", Head.native_eq },
1341 .{ "get_args", Head.get_args },
1342 .{ "get_constructor", Head.get_constructor },
1343 .{ "pbool", Head.pbool },
1344 .{ "get_config", Head.get_config },
1345 .{ "mk_int", Head.mk_int },
1346 .{ "mk_int_weighted", Head.mk_int_weighted },
1347 .{ "int_dist_eq", Head.int_dist_eq },
1348 .{ "print", Head.print_op },
1349 .{ "/.", Head.f_div },
1350 .{ "*.", Head.f_mul },
1351 .{ "+.", Head.f_add },
1352 .{ "-.", Head.f_sub },
1353 .{ "error", Head.error_op },
1354 };
1355
1356 inline for (prims) |prim| {
1357 if (std.mem.eql(u8, name, prim[0])) {
1358 return prim[1];
1359 }
1360 }
1361 return null;
1362 }
1363
1364 fn isLambdaKeyword(token: []const u8) bool {
1365 return std.mem.eql(u8, token, "lam") or
1366 std.mem.eql(u8, token, "lambda") or
1367 std.mem.eql(u8, token, "λ") or
1368 std.mem.eql(u8, token, "fn");
1369 }
1370
1371 fn isIdentifier(token: []const u8) bool {
1372 if (token.len == 0) return false;
1373 const first = token[0];
1374 if (!std.ascii.isAlphabetic(first) and first != '_') return false;
1375 for (token[1..]) |c| {
1376 if (!std.ascii.isAlphanumeric(c) and c != '_') return false;
1377 }
1378 return true;
1379 }
1380
1381 fn isInteger(token: []const u8) bool {
1382 if (token.len == 0) return false;
1383 for (token) |c| {
1384 if (!std.ascii.isDigit(c)) return false;
1385 }
1386 return true;
1387 }
1388
1389 fn isFloat(token: []const u8) bool {
1390 if (token.len == 0) return false;
1391 var has_dot = false;
1392 const start: usize = if (token[0] == '-') 1 else 0;
1393 for (token[start..]) |c| {
1394 if (c == '.') {
1395 if (has_dot) return false;
1396 has_dot = true;
1397 } else if (!std.ascii.isDigit(c)) {
1398 return false;
1399 }
1400 }
1401 return has_dot;
1402 }
1403
1404 pub fn parseExpr(
1405 allocator: Allocator,
1406 source: []const u8,
1407 types: *const TypeRegistry,
1408 defs: *const Definitions,
1409 ) !*PExpr {
1410 const tokens = try tokenize(allocator, source);
1411 defer freeTokens(allocator, tokens);
1412
1413 var parser = Parser.init(allocator, tokens, types, defs);
1414 defer parser.deinit();
1415
1416 const expr = try parser.parseExpr();
1417
1418 if (parser.pos < parser.tokens.len) {
1419 expr.deinit(allocator);
1420 return ParseError.UnexpectedEndOfInput;
1421 }
1422
1423 return expr;
1424 }
1425
1426 test "tokenize basic" {
1427 const allocator = std.testing.allocator;
1428 const tokens = try tokenize(allocator, "(λ x -> x)");
1429 defer freeTokens(allocator, tokens);
1430
1431 try std.testing.expectEqual(@as(usize, 6), tokens.len);
1432 try std.testing.expectEqualStrings("(", tokens[0]);
1433 try std.testing.expectEqualStrings("λ", tokens[1]);
1434 try std.testing.expectEqualStrings("x", tokens[2]);
1435 try std.testing.expectEqualStrings("->", tokens[3]);
1436 try std.testing.expectEqualStrings("x", tokens[4]);
1437 try std.testing.expectEqualStrings(")", tokens[5]);
1438 }
1439
1440 test "tokenize with comment" {
1441 const allocator = std.testing.allocator;
1442 const tokens = try tokenize(allocator, "(x) ;; this is a comment\n(y)");
1443 defer freeTokens(allocator, tokens);
1444
1445 try std.testing.expectEqual(@as(usize, 6), tokens.len);
1446 try std.testing.expectEqualStrings("(", tokens[0]);
1447 try std.testing.expectEqualStrings("x", tokens[1]);
1448 try std.testing.expectEqualStrings(")", tokens[2]);
1449 try std.testing.expectEqualStrings("(", tokens[3]);
1450 try std.testing.expectEqualStrings("y", tokens[4]);
1451 try std.testing.expectEqualStrings(")", tokens[5]);
1452 }
1453
1454 test "tokenize adjacent syntax" {
1455 const allocator = std.testing.allocator;
1456 const tokens = try tokenize(allocator, "a->b=>c|λ[x,y] ;; drop\n`z~q");
1457 defer freeTokens(allocator, tokens);
1458
1459 const expected = [_][]const u8{
1460 "a", "->", "b", "=>", "c", "|",
1461 "λ",
1462 "[", "x", ",", "y", "]", "`",
1463 "z", "~", "q",
1464 };
1465
1466 try std.testing.expectEqual(@as(usize, expected.len), tokens.len);
1467 for (expected, 0..) |token, i| {
1468 try std.testing.expectEqualStrings(token, tokens[i]);
1469 }
1470 }
1471
1472 test "parse identity lambda" {
1473 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1474 defer arena.deinit();
1475 const allocator = arena.allocator();
1476
1477 var types = try TypeRegistry.initWithDefaults(allocator);
1478 var defs = Definitions.init(allocator);
1479
1480 const expr = try parseExpr(allocator, "(λ x -> x)", &types, &defs);
1481
1482 try std.testing.expect(expr.head == .abs);
1483 try std.testing.expectEqualStrings("x", expr.head.abs.var_name);
1484 }
1485
1486 test "parse nat literal" {
1487 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1488 defer arena.deinit();
1489 const allocator = arena.allocator();
1490
1491 var types = try TypeRegistry.initWithDefaults(allocator);
1492 var defs = Definitions.init(allocator);
1493
1494 const expr = try parseExpr(allocator, "3", &types, &defs);
1495
1496 try std.testing.expect(expr.head == .construct);
1497 try std.testing.expectEqualStrings("S", expr.head.construct.constructor);
1498 try std.testing.expectEqual(@as(?i64, 3), expr.maybeConst());
1499 }
1500
1501 test "parse if expression" {
1502 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1503 defer arena.deinit();
1504 const allocator = arena.allocator();
1505
1506 var types = try TypeRegistry.initWithDefaults(allocator);
1507 var defs = Definitions.init(allocator);
1508
1509 const expr = try parseExpr(allocator, "(if true 1 0)", &types, &defs);
1510
1511 try std.testing.expect(expr.head == .case_of);
1512 try std.testing.expectEqual(@as(usize, 2), expr.head.case_of.branches.len);
1513 }
1514
1515 test "parse application" {
1516 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1517 defer arena.deinit();
1518 const allocator = arena.allocator();
1519
1520 var types = try TypeRegistry.initWithDefaults(allocator);
1521 var defs = Definitions.init(allocator);
1522
1523 const expr = try parseExpr(allocator, "((λ x -> x) 42)", &types, &defs);
1524
1525 try std.testing.expect(expr.head == .app);
1526 }
1527
1528 test "parse list literal" {
1529 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1530 defer arena.deinit();
1531 const allocator = arena.allocator();
1532
1533 var types = try TypeRegistry.initWithDefaults(allocator);
1534 var defs = Definitions.init(allocator);
1535
1536 const expr = try parseExpr(allocator, "[1, 2, 3]", &types, &defs);
1537
1538 try std.testing.expect(expr.head == .construct);
1539 try std.testing.expectEqualStrings("Cons", expr.head.construct.constructor);
1540 }
1541
1542 test "parse let expression" {
1543 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1544 defer arena.deinit();
1545 const allocator = arena.allocator();
1546
1547 var types = try TypeRegistry.initWithDefaults(allocator);
1548 var defs = Definitions.init(allocator);
1549
1550 const expr = try parseExpr(allocator, "(let [x 1] x)", &types, &defs);
1551
1552 try std.testing.expect(expr.head == .app);
1553 }
1554
1555 test "parse flip primitive" {
1556 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1557 defer arena.deinit();
1558 const allocator = arena.allocator();
1559
1560 var types = try TypeRegistry.initWithDefaults(allocator);
1561 var defs = Definitions.init(allocator);
1562
1563 const expr = try parseExpr(allocator, "(flip 0.5)", &types, &defs);
1564
1565 try std.testing.expect(expr.head == .flip);
1566 try std.testing.expectEqual(@as(usize, 1), expr.args.len);
1567 }
1568
1569 test "parse factor primitive" {
1570 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1571 defer arena.deinit();
1572 const allocator = arena.allocator();
1573
1574 var types = try TypeRegistry.initWithDefaults(allocator);
1575 var defs = Definitions.init(allocator);
1576
1577 const expr = try parseExpr(allocator, "(factor 0.2)", &types, &defs);
1578
1579 try std.testing.expect(expr.head == .factor);
1580 try std.testing.expectEqual(@as(usize, 1), expr.args.len);
1581 }
1582
1583 test "parse case expression" {
1584 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1585 defer arena.deinit();
1586 const allocator = arena.allocator();
1587
1588 var types = try TypeRegistry.initWithDefaults(allocator);
1589 var defs = Definitions.init(allocator);
1590
1591 const expr = try parseExpr(allocator, "(case true of True => 1 | False => 0)", &types, &defs);
1592
1593 try std.testing.expect(expr.head == .case_of);
1594 try std.testing.expectEqual(@as(usize, 2), expr.head.case_of.branches.len);
1595 try std.testing.expectEqualStrings("True", expr.head.case_of.branches[0].constructor);
1596 try std.testing.expectEqualStrings("False", expr.head.case_of.branches[1].constructor);
1597 }
1598
1599 test "empty uniform rejects" {
1600 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1601 defer arena.deinit();
1602 const allocator = arena.allocator();
1603
1604 var types = try TypeRegistry.initWithDefaults(allocator);
1605 var defs = Definitions.init(allocator);
1606
1607 const result = parseExpr(allocator, "(uniform)", &types, &defs);
1608 try std.testing.expectError(ParseError.InvalidExpression, result);
1609 }
1610
1611 test "empty discrete rejects" {
1612 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1613 defer arena.deinit();
1614 const allocator = arena.allocator();
1615
1616 var types = try TypeRegistry.initWithDefaults(allocator);
1617 var defs = Definitions.init(allocator);
1618
1619 const result = parseExpr(allocator, "(discrete)", &types, &defs);
1620 try std.testing.expectError(ParseError.InvalidExpression, result);
1621 }
1622
1623 test "discrete rejects probabilities not summing to 1" {
1624 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1625 defer arena.deinit();
1626 const allocator = arena.allocator();
1627
1628 var types = try TypeRegistry.initWithDefaults(allocator);
1629 var defs = Definitions.init(allocator);
1630
1631 const result = parseExpr(allocator, "(discrete (True 0.3) (False 0.4))", &types, &defs);
1632 try std.testing.expectError(ParseError.InvalidExpression, result);
1633 }
1634
1635 test "discrete filters zero probabilities" {
1636 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1637 defer arena.deinit();
1638 const allocator = arena.allocator();
1639
1640 var types = try TypeRegistry.initWithDefaults(allocator);
1641 var defs = Definitions.init(allocator);
1642
1643 const result = parseExpr(allocator, "(discrete (True 0.6) (False 0.0) (True 0.4))", &types, &defs);
1644 try std.testing.expect(result != error.InvalidExpression);
1645 }
1646
1647 test "discrete accepts valid distribution" {
1648 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1649 defer arena.deinit();
1650 const allocator = arena.allocator();
1651
1652 var types = try TypeRegistry.initWithDefaults(allocator);
1653 var defs = Definitions.init(allocator);
1654
1655 const result = try parseExpr(allocator, "(discrete (True 0.3) (False 0.7))", &types, &defs);
1656 try std.testing.expect(result.head == .case_of);
1657 }
1658
1659 test "empty case rejects" {
1660 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1661 defer arena.deinit();
1662 const allocator = arena.allocator();
1663
1664 var types = try TypeRegistry.initWithDefaults(allocator);
1665 var defs = Definitions.init(allocator);
1666
1667 const result = parseExpr(allocator, "(case true of )", &types, &defs);
1668 try std.testing.expectError(ParseError.InvalidExpression, result);
1669 }