fuzz.mojo

from emberjson import parse, Array, Object, Value, Null, JSON
from emberjson.utils import write_escaped_string
from std.utils.numerics import isinf
from std.time import monotonic
from std.testing import assert_equal
from std.testing.prop.strategy import Strategy, Rng
from std.testing.prop.strategy.string_strategy import _StringStrategy
from std.testing.prop import PropTest, PropTestConfig
from std.benchmark import keep
from std.time import perf_counter_ns
from std.sys import is_defined
from std.testing import assert_equal


@fieldwise_init
struct JsonStringStrategy(Movable, Strategy):
    comptime Value = String

    def value(self, mut rng: Rng) raises -> Self.Value:
        var j: Value

        if coin_flip(rng):
            j = self.gen_object(rng, 0)
        else:
            j = self.gen_array(rng, 0)

        return String(j)

    def gen_value(self, mut rng: Rng, depth: Int) raises -> Value:
        var max_choice = 7  # 0-7
        if depth > 5:
            max_choice = 5  # 0-5 (Scalars: Null, Int, UInt, Str, Bool, Float)

        var a = rng.rand_int(min=0, max=max_choice)

        if a == 0:
            return Null()
        elif a == 1:
            return rng.rand_scalar[DType.int64]()
        elif a == 2:
            return rng.rand_scalar[DType.uint64]()
        elif a == 3:
            return self.gen_string(rng)
        elif a == 4:
            return coin_flip(rng)
        elif a == 5:
            return rng.rand_scalar[DType.float64]()
        elif a == 6:
            return self.gen_array(rng, depth + 1)
        elif a == 7:
            return self.gen_object(rng, depth + 1)
        else:
            raise Error("Invalid choice")

    def gen_string(self, mut rng: Rng) raises -> String:
        # The stdlib string strategy. The pinned nightly ships the strategy
        # struct but not yet the `String.strategy(...)` extension sugar;
        # switch to `String.strategy(unicode=False, only_printable=True)`
        # once the toolchain resolves it.
        var strat = _StringStrategy(
            min_len=0, max_len=20, unicode=False, only_printable=True
        )
        return strat.value(rng)

    def gen_array(self, mut rng: Rng, depth: Int) raises -> Array:
        var arr = Array()
        var l = rng.rand_int(min=0, max=20 // max(depth, 1))
        arr.reserve(l)
        for _ in range(l):
            arr.append(self.gen_value(rng, depth))
        return arr^

    def gen_object(self, mut rng: Rng, depth: Int) raises -> Object:
        var ob = Object()
        var l = rng.rand_int(min=0, max=20 // max(depth, 1))
        for _ in range(l):
            ob[self.gen_string(rng)] = self.gen_value(rng, depth)
        return ob^


def coin_flip(mut rng: Rng) raises -> Bool:
    return rng.rand_bool()


def main() raises:
    comptime if is_defined["GEN_JSONL"]():
        var rng = Rng(seed=Int(perf_counter_ns()))
        var strat = JsonStringStrategy()

        with open("./bench_data/big_lines_complex.jsonl", "w") as f:
            for _ in range(1_000):
                f.write(strat.value(rng), "\n")

    else:
        print("Running fuzzy tests...")
        var iters = 100

        @parameter
        def test_parse(s: String) raises:
            var rng = Rng(seed=Int(perf_counter_ns()))
            var j: Value = {}
            if iters % 4 == 0:
                var start = rng.rand_int(min=0, max=s.byte_length())
                var end = rng.rand_int(min=start, max=s.byte_length())
                var corrupted = s[byte=start:end]
                try:
                    j = parse(corrupted)
                except:
                    # Main thing is we don't want this to crash.
                    # But don't enforce failure on the off chance this slicing happens to
                    # produce valid json.
                    pass
            else:
                j = parse(s)
                assert_equal(String(j), s)
            iters -= 1
            keep(j)

        var test = PropTest(config=PropTestConfig(runs=iters))
        test.test[test_parse](JsonStringStrategy())
        print("Test passed!")