# This file is a part of Julia. License is MIT: https://julialang.org/license

module HeapSnapshot

# SoA layout to reduce padding
struct Edges
    type::Vector{UInt32}       # index into `snapshot.meta.edge_types`
    name_or_index::Vector{UInt} # Either an index into `snapshot.strings`, or the index in an array, depending on edge_type
    to_pos::Vector{UInt}   # index into `snapshot.nodes`
end
function Edges(n::Int)
    Edges(
        Vector{UInt32}(undef, n),
        Vector{UInt}(undef, n),
        Vector{UInt}(undef, n),
    )
end
Base.length(n::Edges) = length(n.type)

# trace_node_id and detachedness are always 0 in the snapshots Julia produces so we don't store them
struct Nodes
    type::Vector{UInt32}         # index into `snapshot.meta.node_types`
    name_idx::Vector{UInt32} # index into `snapshot.strings`
    id::Vector{UInt}           # unique id, in julia it is the address of the object
    self_size::Vector{Int}     # size of the object itself, not including the size of its fields
    edge_count::Vector{UInt}   # number of outgoing edges
    # This is the main complexity of the .heapsnapshot format, and it's the reason we need
    # to read in all the data before writing it out. The edges vector contains all edges,
    # but organized by which node they came from. First, it contains all the edges coming
    # out of node 0, then all edges leaving node 1, etc. So we store the edges in that
    # order, which means an edge's position is determined by the number of edges leaving
    # all the preceding nodes, and no per-node index is needed.
    edges::Edges               # outgoing edges, grouped by source node
end
function Nodes(n::Int, e::Int)
    Nodes(
        Vector{UInt32}(undef, n),
        Vector{UInt32}(undef, n),
        Vector{UInt}(undef, n),
        Vector{Int}(undef, n),
        Vector{UInt}(undef, n),
        Edges(e),
    )
end
Base.length(n::Nodes) = length(n.type)

const k_node_number_of_fields = 7

# the streaming .edges format writes each edge as a packed
# (type, name_or_index, from_node, to_node) record; see `serialize_edge` in
# src/gc-heap-snapshot.c
const k_edge_record_bytes = sizeof(UInt32) + 3 * sizeof(UInt)
const k_edge_from_node_offset = sizeof(UInt32) + sizeof(UInt)

# Like Base.dec, but doesn't allocate a string and writes directly to the io object
# We know all of the numbers we're about to write fit into a UInt and are non-negative
let _dec_d100 = UInt16[(0x30 + i % 10) << 0x8 + (0x30 + i ÷ 10) for i = 0:99]
    global _write_decimal_number
    _write_decimal_number(io, x::Integer, buf) = _write_decimal_number(io, unsigned(x), buf)
    function _write_decimal_number(io, x::Unsigned, digits_buf)
        buf = digits_buf
        n = ndigits(x)
        i = n
        @inbounds while i >= 2
            d, r = divrem(x, 0x64)
            d100 = _dec_d100[(r % Int)::Int + 1]
            buf[i-1] = d100 % UInt8
            buf[i] = (d100 >> 0x8) % UInt8
            x = oftype(x, d)
            i -= 2
        end
        if i > 0
            @inbounds buf[i] = 0x30 + (rem(x, 0xa) % UInt8)::UInt8
        end
        write(io, @view buf[max(i, 1):n])
    end
end

const k_output_chunk_bytes = 1 << 20

# Hand the buffered output over to `io` once at least `limit` bytes have accumulated,
# reusing the buffer. Pass `limit = 0` to flush what's left.
function _maybe_flush(io::IO, buf::IOBuffer, limit::Int = k_output_chunk_bytes)
    n = position(buf)
    if n >= limit
        GC.@preserve buf unsafe_write(io, pointer(buf.data), n)
        truncate(buf, 0)
        seekstart(buf)
    end
    return nothing
end

"""
    assemble_snapshot(in_prefix::AbstractString, out_file::AbstractString = in_prefix)

Assemble a `.heapsnapshot` file from the four files produced by
`Profile.take_heap_snapshot(...; streaming=true)` with prefix `in_prefix`.
"""
function assemble_snapshot(in_prefix, out_file::AbstractString = in_prefix)
    # assemble into a temporary sibling and only replace `out_file` on success, so
    # that a failure doesn't leave a partial snapshot or destroy a pre-existing one
    tmp_file = tempname(dirname(abspath(out_file)); cleanup=false)
    try
        open(tmp_file, "w") do io
            assemble_snapshot(in_prefix, io)
        end
        # `rename` replaces atomically, unlike `mv(...; force=true)` which unlinks first
        Base.Filesystem.rename(tmp_file, out_file)
    finally
        rm(tmp_file; force=true)
    end
    return nothing
end

# Manually parse and write the .json files, given that we don't have JSON import/export in
# julia's stdlibs.
function assemble_snapshot(in_prefix, io::IO)
    preamble = read(string(in_prefix, ".metadata.json"), String)
    pos = last(findfirst("node_count\":", preamble)) + 1
    endpos = findnext(==(','), preamble, pos) - 1
    node_count = parse(Int, String(@view preamble[pos:endpos]))

    pos = last(findnext("edge_count\":", preamble, endpos)) + 1
    endpos = findnext(==(','), preamble, pos) - 1
    edge_count = parse(Int, String(@view preamble[pos:endpos]))

    nodes = Nodes(node_count, edge_count)

    # whether each node has at least one incoming edge; the uber node has none by
    # construction, so mark it up front
    has_incoming = falses(node_count)
    node_count > 0 && (has_incoming[1] = true)
    # Parse nodes with empty edge counts that we need to fill later.
    # Read each part file into memory up front: parsing field-at-a-time from an
    # IOStream is dominated by per-call locking overhead.
    let nodes_file = IOBuffer(read(string(in_prefix, ".nodes")))
        for i in 1:length(nodes)
            node_type = read(nodes_file, UInt32)
            node_name_idx = read(nodes_file, UInt)
            id = read(nodes_file, UInt)
            self_size = read(nodes_file, Int)
            trace_node_id = read(nodes_file, Int) # always 0 in snapshots Julia produces
            trace_node_id == 0 ||
                error("malformed nodes file `$(in_prefix).nodes`: nonzero trace_node_id $(trace_node_id) for node $(i)")
            detachedness = read(nodes_file, Int8) # always 0 in snapshots Julia produces
            detachedness == 0 ||
                error("malformed nodes file `$(in_prefix).nodes`: nonzero detachedness $(detachedness) for node $(i)")

            nodes.type[i] = node_type
            nodes.name_idx[i] = node_name_idx
            nodes.id[i] = id
            nodes.self_size[i] = self_size
            nodes.edge_count[i] = 0 # edge_count
        end
    end

    # Parse the edges to fill in the edge counts for nodes and correct the to_node offsets.
    # The edges are stored grouped by source node, which is the order they have to be
    # emitted in, so a first pass counts the edges leaving each node and a second one
    # places each edge at its position within its source node's group.
    let edges_data = read(string(in_prefix, ".edges"))
        n_edges = length(nodes.edges)
        length(edges_data) == n_edges * k_edge_record_bytes ||
            error("malformed edges file `$(in_prefix).edges`: expected $(n_edges * k_edge_record_bytes) bytes for $(n_edges) edges, got $(length(edges_data))")
        edges_file = IOBuffer(edges_data)
        for i in 1:n_edges
            seek(edges_file, (i - 1) * k_edge_record_bytes + k_edge_from_node_offset)
            from_node = read(edges_file, UInt)
            from_node < node_count ||
                error("malformed edges file `$(in_prefix).edges`: from_node $(from_node) is out of range for $(node_count) nodes")
            nodes.edge_count[from_node + 1] += UInt(1)  # C and JSON use 0-based indexing
        end
        # position in `nodes.edges` of the next edge leaving each node
        cursor = Vector{Int}(undef, length(nodes))
        pos = 1
        for n in 1:length(nodes)
            cursor[n] = pos
            pos += nodes.edge_count[n]
        end
        seekstart(edges_file)
        for _ in 1:n_edges
            edge_type = read(edges_file, UInt32)
            edge_name_or_index = read(edges_file, UInt)
            from_node = read(edges_file, UInt)
            to_node = read(edges_file, UInt)

            i = cursor[from_node + 1]
            cursor[from_node + 1] = i + 1
            nodes.edges.type[i] = edge_type
            nodes.edges.name_or_index[i] = edge_name_or_index
            nodes.edges.to_pos[i] = to_node * k_node_number_of_fields # 7 fields per node, the streaming format doesn't multiply the offset by 7
            to_node < node_count ||
                error("malformed edges file `$(in_prefix).edges`: to_node $(to_node) is out of range for $(node_count) nodes")
            has_incoming[to_node + 1] = true
        end
    end

    if !all(has_incoming)
        orphans = findall(!, has_incoming) .- 1 # C and JSON use 0-based indexing
        error("malformed snapshot `$(in_prefix)`: $(length(orphans)) of $(length(nodes)) nodes have no incoming edges: ",
              join(Iterators.take(orphans, 10), ", "), length(orphans) > 10 ? ", ..." : "")
    end

    _digits_buf = zeros(UInt8, ndigits(typemax(UInt)))
    # the snapshot is emitted a field at a time, and a write to an `IOStream` takes a
    # lock and enters the runtime, so batch the output and hand it to `io` in chunks
    out = IOBuffer(sizehint = k_output_chunk_bytes)
    println(out, @view(preamble[1:end-1]), ",") # remove trailing "}" to reopen the object

    println(out, "\"nodes\":[")
    for i in 1:length(nodes)
        i > 1 && println(out, ",")
        _write_decimal_number(out, nodes.type[i], _digits_buf)
        print(out, ",")
        _write_decimal_number(out, nodes.name_idx[i], _digits_buf)
        print(out, ",")
        _write_decimal_number(out, nodes.id[i], _digits_buf)
        print(out, ",")
        _write_decimal_number(out, nodes.self_size[i], _digits_buf)
        print(out, ",")
        _write_decimal_number(out, nodes.edge_count[i], _digits_buf)
        print(out, ",0,0")
        _maybe_flush(io, out)
    end
    print(out, "],\n")
    print(out, "\"edges\":[")
    i = 1
    for n in 1:length(nodes)
        for _ in 1:nodes.edge_count[n]
            i > 1 && print(out, ",")
            println(out)
            _write_decimal_number(out, nodes.edges.type[i], _digits_buf)
            print(out, ",")
            _write_decimal_number(out, nodes.edges.name_or_index[i], _digits_buf)
            print(out, ",")
            _write_decimal_number(out, nodes.edges.to_pos[i], _digits_buf)
            if !(nodes.edges.to_pos[i] % k_node_number_of_fields == 0)
                @warn "Bug in to_pos for edge $i from node $n: $(nodes.edges.to_pos[i])"
            end
            i += 1
            _maybe_flush(io, out)
        end
    end
    println(out, "],")

    # not used. Required by microsoft/vscode-v8-heap-tools
    # This order of these fields is required by chrome dev tools otherwise loading fails
    println(out, "\"trace_function_infos\":[],")
    println(out, "\"trace_tree\":[],")
    println(out, "\"samples\":[],")
    println(out, "\"locations\":[],")

    println(out, "\"strings\":[")
    let strings_io = IOBuffer(read(string(in_prefix, ".strings")))
        first = true
        while !eof(strings_io)
            str_size = read(strings_io, UInt)
            str_bytes = read(strings_io, str_size)
            length(str_bytes) == str_size ||
                error("truncated strings file: `$(in_prefix).strings` (expected $(str_size) bytes, got $(length(str_bytes)))")
            str = String(str_bytes)
            if first
                first = false
            else
                print(out, ",\n")
            end
            print_str_escape_json(out, str)
            _maybe_flush(io, out)
        end
    end
    print(out, "]}")
    _maybe_flush(io, out, 0)

    return nothing
end

"""
    cleanup_streamed_files(prefix::AbstractString)

Remove files streamed during `take_heap_snapshot` in streaming mode.
"""
function cleanup_streamed_files(prefix::AbstractString)
    # tolerate missing files, e.g. from a failure partway through streaming
    rm(string(prefix, ".metadata.json"); force=true)
    rm(string(prefix, ".nodes"); force=true)
    rm(string(prefix, ".edges"); force=true)
    rm(string(prefix, ".strings"); force=true)
    return nothing
end

function print_str_escape_json(stream::IO, s::AbstractString)
    print(stream, '"')
    # write the string in runs, so that a stretch of characters needing no escaping
    # costs one write rather than one per character
    runstart = firstindex(s)
    for (i, c) in pairs(s)
        escaped = if c == '"'
            "\\\""
        elseif c == '\\'
            "\\\\"
        elseif c == '\b'
            "\\b"
        elseif c == '\f'
            "\\f"
        elseif c == '\n'
            "\\n"
        elseif c == '\r'
            "\\r"
        elseif c == '\t'
            "\\t"
        elseif '\x00' <= c <= '\x1f'
            string("\\u", lpad(string(UInt16(c), base=16), 4, '0'))
        elseif !isvalid(c)
            # we have to do this because vscode's viewer doesn't like the replace character
            "[invalid unicode character]"
        else
            continue
        end
        i > runstart && print(stream, SubString(s, runstart, prevind(s, i)))
        print(stream, escaped)
        runstart = nextind(s, i)
    end
    runstart <= lastindex(s) && print(stream, SubString(s, runstart))
    print(stream, '"')
end

end
