module CallGraphs

import Base: hash, isless
export FileFunction, combine, emacs_highlighting, findcallers, findgc, highlightgc, parsedot, parsedots

const rnode = r"Node(\w+)"     # for extracting the Node# from a .dot file
const rlabel = r"\"{(.*)}\""   # for extracting the function name from a "label" in a .dot file
const rfname = r"(\w+)\(.*\)"  # for stripping the signature from a function
const ralphanum = r"\w+"       # for extracting chunks of alphanumeric characters (e.g., function names)

type FileFunction
    file::ASCIIString
    func::ASCIIString
end

isless(cf1::FileFunction, cf2::FileFunction) = cf1.file < cf2.file ? true : (cf1.file == cf2.file && cf1.func < cf2.func)
==(cf1::FileFunction, cf2::FileFunction) = cf1.file == cf2.file && cf1.func == cf2.func
hash(cf::FileFunction, h::UInt) = hash(cf.file, hash(cf.func, h))

### Parsing .dot callgraph files
@doc """
`cgs = parsedots(dirname=".")` parses all callgraph `*.dot` files in the specified directory.
See `combine` for merging them into a single call graph, and `parsedot` for detailed information
about the encoding within `cgs`.
""" ->
function parsedots(dirname=".")
    isdir(dirname) || error("$dirname is not a directory")
    fls = readdir(dirname)
    cgs = Array(Any, 0)
    for fl in fls
        basename, ext = splitext(fl)
        ext == ".dot" || continue
        cg = parsedot(joinpath(dirname, fl))
        push!(cgs, (basename, cg))
    end
    cgs
end


@doc """
`nodes, labels, calls, calledby, externally_defined = parsedot(filename)`

Parses a single callgraph file, stored in GraphViz `*.dot` file, as generated by `llvm`'s
`opt -analyze -dot-callgraph`. Nodes in the graph (functions) are represented as `UInt32`
values, most of which have a `label` (the name, and sometimes the signature, of the function).
`nodes` is a dictionary convering `lbl=>node`, and `labels` performs
the opposite mapping. For a function represented by `node`, `calls[node]` returns the set
of nodes that `node` calls. Conversely, `calledby[node]` returns the list of functions
that call `node`.  Finally, some of the functions lack definitions in the current file, and
are all represented as calling the node `externally_defined`.

This is a low-level interface; usually you want to use `parsedots` and `combine`.
""" ->
function parsedot(io::IO)
    lines = split(readall(io), '\n')
    labels = Dict{UInt32,ASCIIString}()
    nodes = Dict{ASCIIString,UInt32}()
    calls = Dict{UInt32,Set{UInt32}}()
    calledby = Dict{UInt32,Set{UInt32}}()
    local external_caller
    for line in lines
        mnode = match(rnode, line)
        mnode == nothing && continue
        node = UInt32(parseint(mnode.captures[1]))
        if contains(line, "label")
            # A new node
            mlabel = match(rlabel, line)
            label = mlabel.captures[1]
            if label == "external node"
                external_caller = node
                continue
            end
            labels[node] = label
            nodes[label] = node
            calls[node] = Set{UInt32}()
            if !haskey(calledby, node)
                calledby[node] = Set{UInt32}()
            end
        else
            node == external_caller && continue
            mnode = match(rnode, line, mnode.offsets[1])
            mnode == nothing && error("Don't understand format of ", line)
            isempty(mnode.captures) && error("Don't understand format of ", line)
            callee = UInt32(parseint(mnode.captures[1]))
            push!(calls[node], callee)
            if haskey(calledby, callee)
                push!(calledby[callee], node)
            else
                calledby[callee] = Set{UInt32}(node)
            end
        end
    end
    for (key, val) in calledby
        delete!(val, external_caller)
    end
    # Find a node with no label---this corresponds to externally-defined functions
    local externally_defined
    for (nd, s) in calledby
        if !haskey(labels, nd)
            externally_defined = nd
            break
        end
    end
    nodes, labels, calls, calledby, externally_defined
end

function parsedot(filename::AbstractString)
    results = open(filename) do file
        parsedot(file)
    end
    results
end

### Combining callgraphs from many files into a single callgraph
@doc """
`calls, calledby = combine(cgs...)` merges the callgraphs from multiple files into a single
large callgraph. Both `calls` and `calledby` are dictionaries, where the `key` is
a `FileFunction`, a type which encodes the `file` and `func` as strings.

`calls[FileFunction("file1.c", "function8")]` will return the set of functions that `"function8"` calls.
`calledby[FileFunction("file1.c", "function8")]` will return the set of functions calling `"function8"`.

In cases where `opt -analyze -dot-callgraph` retained the function argument signature,
`combine` will discard it.

See also `parsedots`, which generates `cgs`.
""" ->
function combine(cgs...)
    functions = Array(FileFunction, 0)
    lookup = Dict{ASCIIString, FileFunction}()
    for cg in cgs
        filename, results = cg
        nodes, labels, calls, calledby, externally_defined = results
        for (nd,lbl) in labels
            olbl = lbl
            # Discard argument signatures and just keep the function name
            m = match(rfname, lbl)
            if m != nothing
                lbl = m.captures[1]
            end
            func = FileFunction(filename, lbl)
            push!(functions, func)
            lookup[olbl] = func
        end
    end
    calls = Dict{FileFunction, Set{FileFunction}}()
    calledby = Dict{FileFunction, Set{FileFunction}}()
    for cg in cgs
        filename, results = cg
        nodes, labels, calls1, calledby1, externally_defined = results
        for (nd,lbl) in labels
            func = lookup[lbl]
            if !haskey(calls, func)
                calls[func] = Set{FileFunction}()
            end
            cls = calls[func]
            for c in calls1[nd]
                c == externally_defined && continue
                callee = lookup[labels[c]]
                push!(cls, callee)
                if !haskey(calledby, callee)
                    calledby[callee] = Set{FileFunction}()
                end
                push!(calledby[callee], func)
            end
        end
    end
    calls, calledby
end

### Finding the list of functions that call funcname
@doc """
`callers = findcallers(calledby, funcname; indirect=true)` returns a set of `FileFunction`s
that call the function with name `funcname`. `calledby` is the dictionary
returned by `combine`. If `indirect` is true, then `callers` includes functions that might
call `funcname` even indirectly, through other functions.
""" ->
function findcallers(calledby, funcname; indirect::Bool=true)
    local funckey
    found = false
    for (key, callers) in calledby
        if key.func == funcname
            found = true
            funckey = key
        end
    end
    found || error("$funcname not found")
    callers = calledby[funckey]
    push!(callers, funckey)
    # Handle recursive/indirect calls by iterating the call chain until
    # we reach steady-state
    if indirect
        while true
            newcallers = Set{FileFunction}()
            for func in callers
                push!(newcallers, func)
                if haskey(calledby, func)
                    union!(newcallers, calledby[func])
                end
            end
            if isempty(setdiff(newcallers, callers))
                break
            end
            callers = newcallers
        end
    end
    callers
end

@doc """
`names = uniquenames(filefuncs)` takes a vector of `FileFunction`s and returns
a set of unique function names.
""" ->
function uniquenames(filefuncs::Vector{FileFunction})
    names = Set{ASCIIString}()
    for f in filefuncs
        push!(names, f.func)
    end
    names
end

include("gc.jl")

end
