|
| 1 | +# Reference: https://github.com/probcomp/Gen.jl/blob/91d798f2d2f0c175b1be3dc6daf3a10a8acf5da3/src/choice_map.jl#L104 |
| 2 | + |
| 3 | +function _format_array(arr::AbstractArray; n_show::Int=3, indent::Int=0) |
| 4 | + nd = ndims(arr) |
| 5 | + if nd == 0 |
| 6 | + return string(arr[]) |
| 7 | + elseif nd == 1 |
| 8 | + len = length(arr) |
| 9 | + if len <= 2 * n_show |
| 10 | + return "[" * join(arr, " ") * "]" |
| 11 | + end |
| 12 | + first_part = join(arr[1:n_show], " ") |
| 13 | + last_part = join(arr[(end - n_show + 1):end], " ") |
| 14 | + return "[$first_part ... $last_part]" |
| 15 | + else |
| 16 | + n_slices = size(arr, 1) |
| 17 | + indent_str = " "^(indent + 1) |
| 18 | + |
| 19 | + if n_slices <= 2 * n_show |
| 20 | + slice_strs = [ |
| 21 | + _format_array(selectdim(arr, 1, i); n_show=n_show, indent=indent + 1) for |
| 22 | + i in 1:n_slices |
| 23 | + ] |
| 24 | + return "[" * join(slice_strs, "\n" * indent_str) * "]" |
| 25 | + else |
| 26 | + first_slices = [ |
| 27 | + _format_array(selectdim(arr, 1, i); n_show=n_show, indent=indent + 1) for |
| 28 | + i in 1:n_show |
| 29 | + ] |
| 30 | + last_slices = [ |
| 31 | + _format_array(selectdim(arr, 1, i); n_show=n_show, indent=indent + 1) for |
| 32 | + i in (n_slices - n_show + 1):n_slices |
| 33 | + ] |
| 34 | + return "[" * |
| 35 | + join(first_slices, "\n" * indent_str) * |
| 36 | + "\n" * |
| 37 | + indent_str * |
| 38 | + "..." * |
| 39 | + "\n" * |
| 40 | + indent_str * |
| 41 | + join(last_slices, "\n" * indent_str) * |
| 42 | + "]" |
| 43 | + end |
| 44 | + end |
| 45 | +end |
| 46 | + |
| 47 | +function _format_digest(value; n_show::Int=3) |
| 48 | + if isa(value, Tuple) |
| 49 | + if length(value) == 1 |
| 50 | + return _format_digest(value[1]; n_show=n_show) |
| 51 | + else |
| 52 | + formatted = [_format_digest(v; n_show=n_show) for v in value] |
| 53 | + return "(" * join(formatted, ", ") * ")" |
| 54 | + end |
| 55 | + elseif isa(value, AbstractArray) |
| 56 | + return _format_array(value; n_show=n_show, indent=0) |
| 57 | + else |
| 58 | + return string(value) |
| 59 | + end |
| 60 | +end |
| 61 | + |
| 62 | +function _show_pretty(io::IO, trace::Trace, pre::Int, vert_bars::Tuple) |
| 63 | + VERT = '\u2502' |
| 64 | + PLUS = '\u251C' |
| 65 | + HORZ = '\u2500' |
| 66 | + LAST = '\u2514' |
| 67 | + |
| 68 | + indent_vert = vcat(Char[' ' for _ in 1:pre], Char[VERT, '\n']) |
| 69 | + indent = vcat(Char[' ' for _ in 1:pre], Char[PLUS, HORZ, HORZ, ' ']) |
| 70 | + indent_last = vcat(Char[' ' for _ in 1:pre], Char[LAST, HORZ, HORZ, ' ']) |
| 71 | + |
| 72 | + for i in vert_bars |
| 73 | + indent_vert[i] = VERT |
| 74 | + indent[i] = VERT |
| 75 | + indent_last[i] = VERT |
| 76 | + end |
| 77 | + |
| 78 | + indent_vert_str = join(indent_vert) |
| 79 | + indent_str = join(indent) |
| 80 | + indent_last_str = join(indent_last) |
| 81 | + |
| 82 | + sorted_choices = sort(collect(trace.choices); by=x -> x[1]) |
| 83 | + n = length(sorted_choices) |
| 84 | + |
| 85 | + if trace.retval !== nothing |
| 86 | + n += 1 |
| 87 | + end |
| 88 | + |
| 89 | + if trace.weight !== nothing |
| 90 | + n += 1 |
| 91 | + end |
| 92 | + |
| 93 | + cur = 1 |
| 94 | + |
| 95 | + if trace.retval !== nothing |
| 96 | + print(io, indent_vert_str) |
| 97 | + retval_str = _format_digest(trace.retval) |
| 98 | + print(io, (cur == n ? indent_last_str : indent_str) * "retval : $retval_str\n") |
| 99 | + cur += 1 |
| 100 | + end |
| 101 | + |
| 102 | + if trace.weight !== nothing |
| 103 | + print(io, indent_vert_str) |
| 104 | + print(io, (cur == n ? indent_last_str : indent_str) * "weight : $(trace.weight)\n") |
| 105 | + cur += 1 |
| 106 | + end |
| 107 | + |
| 108 | + for (key, value) in sorted_choices |
| 109 | + print(io, indent_vert_str) |
| 110 | + value_str = _format_digest(value) |
| 111 | + if contains(value_str, '\n') |
| 112 | + indent_continuation = " "^(length(indent_str) + length(repr(key)) + 3) |
| 113 | + value_str = replace(value_str, "\n" => "\n" * indent_continuation) |
| 114 | + end |
| 115 | + print(io, (cur == n ? indent_last_str : indent_str) * "$(repr(key)) : $value_str\n") |
| 116 | + cur += 1 |
| 117 | + end |
| 118 | + |
| 119 | + sorted_subtraces = sort(collect(trace.subtraces); by=x -> x[1]) |
| 120 | + n += length(sorted_subtraces) |
| 121 | + |
| 122 | + for (key, subtrace) in sorted_subtraces |
| 123 | + print(io, indent_vert_str) |
| 124 | + print(io, (cur == n ? indent_last_str : indent_str) * "subtrace on $(repr(key))\n") |
| 125 | + _show_pretty( |
| 126 | + io, subtrace, pre + 4, cur == n ? (vert_bars...,) : (vert_bars..., pre + 1) |
| 127 | + ) |
| 128 | + cur += 1 |
| 129 | + end |
| 130 | +end |
| 131 | + |
| 132 | +function Base.show(io::IO, ::MIME"text/plain", trace::Trace) |
| 133 | + println(io, "Trace:") |
| 134 | + if isempty(trace.choices) && trace.retval === nothing && trace.weight === nothing |
| 135 | + println(io, " (empty)") |
| 136 | + else |
| 137 | + _show_pretty(io, trace, 0, ()) |
| 138 | + end |
| 139 | +end |
| 140 | + |
| 141 | +function Base.show(io::IO, trace::Trace) |
| 142 | + if get(io, :compact, false) |
| 143 | + choices_count = length(trace.choices) |
| 144 | + has_retval = trace.retval !== nothing |
| 145 | + print(io, "Trace($(choices_count) choices") |
| 146 | + if has_retval |
| 147 | + print(io, ", retval=$(trace.retval), weight=$(trace.weight)") |
| 148 | + end |
| 149 | + print(io, ")") |
| 150 | + else |
| 151 | + show(io, MIME"text/plain"(), trace) |
| 152 | + end |
| 153 | +end |
0 commit comments