/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "builtin/TestingFunctions.h"
#include "js/UbiNode.h"
#include "js/UbiNodeDominatorTree.h"
#include "js/UbiNodePostOrder.h"
#include "jsapi-tests/tests.h"
#include "vm/SavedFrame.h"

using JS::RootedObject;
using JS::RootedScript;
using JS::RootedString;
using namespace js;

// A helper JS::ubi::Node concrete implementation that can be used to make mock
// graphs for testing traversals with.
struct FakeNode
{
    char                name;
    JS::ubi::EdgeVector edges;

    explicit FakeNode(char name) : name(name), edges() { }

    bool addEdgeTo(FakeNode& referent) {
        JS::ubi::Node node(&referent);
        return edges.emplaceBack(nullptr, node);
    }
};

namespace JS {
namespace ubi {

template<>
struct Concrete<FakeNode> : public Base
{
    static const char16_t concreteTypeName[];
    const char16_t* typeName() const override { return concreteTypeName; }

    UniquePtr<EdgeRange> edges(JSRuntime* rt, bool wantNames) const override {
        return UniquePtr<EdgeRange>(js_new<PreComputedEdgeRange>(get().edges));
    }

    Node::Size size(mozilla::MallocSizeOf) const override {
        return 1;
    }

    static void construct(void* storage, FakeNode* ptr) { new (storage) Concrete(ptr); }

  protected:
    explicit Concrete(FakeNode* ptr) : Base(ptr) { }
    FakeNode& get() const { return *static_cast<FakeNode*>(ptr); }
};

const char16_t Concrete<FakeNode>::concreteTypeName[] = MOZ_UTF16("FakeNode");

} // namespace ubi
} // namespace JS

// ubi::Node::zone works
BEGIN_TEST(test_ubiNodeZone)
{
    RootedObject global1(cx, JS::CurrentGlobalOrNull(cx));
    CHECK(global1);
    CHECK(JS::ubi::Node(global1).zone() == cx->zone());

    RootedObject global2(cx, JS_NewGlobalObject(cx, getGlobalClass(), nullptr,
                                                JS::FireOnNewGlobalHook));
    CHECK(global2);
    CHECK(global1->zone() != global2->zone());
    CHECK(JS::ubi::Node(global2).zone() == global2->zone());
    CHECK(JS::ubi::Node(global2).zone() != global1->zone());

    JS::CompileOptions options(cx);

    // Create a string and a script in the original zone...
    RootedString string1(cx, JS_NewStringCopyZ(cx, "Simpson's Individual Stringettes!"));
    CHECK(string1);
    RootedScript script1(cx);
    CHECK(JS::Compile(cx, options, "", 0, &script1));

    {
        // ... and then enter global2's zone and create a string and script
        // there, too.
        JSAutoCompartment ac(cx, global2);

        RootedString string2(cx, JS_NewStringCopyZ(cx, "A million household uses!"));
        CHECK(string2);
        RootedScript script2(cx);
        CHECK(JS::Compile(cx, options, "", 0, &script2));

        CHECK(JS::ubi::Node(string1).zone() == global1->zone());
        CHECK(JS::ubi::Node(script1).zone() == global1->zone());

        CHECK(JS::ubi::Node(string2).zone() == global2->zone());
        CHECK(JS::ubi::Node(script2).zone() == global2->zone());
    }

    return true;
}
END_TEST(test_ubiNodeZone)

// ubi::Node::compartment works
BEGIN_TEST(test_ubiNodeCompartment)
{
    RootedObject global1(cx, JS::CurrentGlobalOrNull(cx));
    CHECK(global1);
    CHECK(JS::ubi::Node(global1).compartment() == cx->compartment());

    RootedObject global2(cx, JS_NewGlobalObject(cx, getGlobalClass(), nullptr,
                                                JS::FireOnNewGlobalHook));
    CHECK(global2);
    CHECK(global1->compartment() != global2->compartment());
    CHECK(JS::ubi::Node(global2).compartment() == global2->compartment());
    CHECK(JS::ubi::Node(global2).compartment() != global1->compartment());

    JS::CompileOptions options(cx);

    // Create a script in the original compartment...
    RootedScript script1(cx);
    CHECK(JS::Compile(cx, options, "", 0, &script1));

    {
        // ... and then enter global2's compartment and create a script
        // there, too.
        JSAutoCompartment ac(cx, global2);

        RootedScript script2(cx);
        CHECK(JS::Compile(cx, options, "", 0, &script2));

        CHECK(JS::ubi::Node(script1).compartment() == global1->compartment());
        CHECK(JS::ubi::Node(script2).compartment() == global2->compartment());
    }

    return true;
}
END_TEST(test_ubiNodeCompartment)

BEGIN_TEST(test_ubiNodeJSObjectConstructorName)
{
    JS::RootedValue val(cx);
    EVAL("this.Ctor = function Ctor() {}; new Ctor", &val);
    CHECK(val.isObject());

    mozilla::UniquePtr<char16_t[], JS::FreePolicy> ctorName;
    CHECK(JS::ubi::Node(&val.toObject()).jsObjectConstructorName(cx, ctorName));
    CHECK(ctorName);
    CHECK(js_strcmp(ctorName.get(), MOZ_UTF16("Ctor")) == 0);

    return true;
}
END_TEST(test_ubiNodeJSObjectConstructorName)

template <typename F, typename G>
static bool
checkString(const char* expected, F fillBufferFunction, G stringGetterFunction)
{
    auto expectedLength = strlen(expected);
    char16_t buf[1024];
    if (fillBufferFunction(mozilla::RangedPtr<char16_t>(buf, 1024), 1024) != expectedLength ||
        !EqualChars(buf, expected, expectedLength))
    {
        return false;
    }

    auto string = stringGetterFunction();
    // Expecting a |JSAtom*| from a live |JS::ubi::StackFrame|.
    if (!string.template is<JSAtom*>() ||
        !StringEqualsAscii(string.template as<JSAtom*>(), expected))
    {
        return false;
    }

    return true;
}

BEGIN_TEST(test_ubiStackFrame)
{
    CHECK(js::DefineTestingFunctions(cx, global, false, false));

    JS::RootedValue val(cx);
    CHECK(evaluate("(function one() {                      \n"  // 1
                   "  return (function two() {             \n"  // 2
                   "    return (function three() {         \n"  // 3
                   "      return saveStack();              \n"  // 4
                   "    }());                              \n"  // 5
                   "  }());                                \n"  // 6
                   "}());                                  \n", // 7
                   "filename.js",
                   1,
                   &val));

    CHECK(val.isObject());
    JS::RootedObject obj(cx, &val.toObject());

    CHECK(obj->is<SavedFrame>());
    JS::Rooted<SavedFrame*> savedFrame(cx, &obj->as<SavedFrame>());

    JS::ubi::StackFrame ubiFrame(savedFrame);

    // All frames should be from the "filename.js" source.
    while (ubiFrame) {
        CHECK(checkString("filename.js",
                          [&] (mozilla::RangedPtr<char16_t> ptr, size_t length) {
                              return ubiFrame.source(ptr, length);
                          },
                          [&] {
                              return ubiFrame.source();
                          }));
        ubiFrame = ubiFrame.parent();
    }

    ubiFrame = savedFrame;

    auto bufferFunctionDisplayName = [&] (mozilla::RangedPtr<char16_t> ptr, size_t length) {
        return ubiFrame.functionDisplayName(ptr, length);
    };
    auto getFunctionDisplayName = [&] {
        return ubiFrame.functionDisplayName();
    };

    CHECK(checkString("three", bufferFunctionDisplayName, getFunctionDisplayName));
    CHECK(ubiFrame.line() == 4);

    ubiFrame = ubiFrame.parent();
    CHECK(checkString("two", bufferFunctionDisplayName, getFunctionDisplayName));
    CHECK(ubiFrame.line() == 3);

    ubiFrame = ubiFrame.parent();
    CHECK(checkString("one", bufferFunctionDisplayName, getFunctionDisplayName));
    CHECK(ubiFrame.line() == 2);

    ubiFrame = ubiFrame.parent();
    CHECK(ubiFrame.functionDisplayName().is<JSAtom*>());
    CHECK(ubiFrame.functionDisplayName().as<JSAtom*>() == nullptr);
    CHECK(ubiFrame.line() == 1);

    ubiFrame = ubiFrame.parent();
    CHECK(!ubiFrame);

    return true;
}
END_TEST(test_ubiStackFrame)

BEGIN_TEST(test_ubiCoarseType)
{
    // Test that our explicit coarseType() overrides work as expected.

    JSObject* obj = nullptr;
    CHECK(JS::ubi::Node(obj).coarseType() == JS::ubi::CoarseType::Object);

    JSScript* script = nullptr;
    CHECK(JS::ubi::Node(script).coarseType() == JS::ubi::CoarseType::Script);

    js::LazyScript* lazyScript = nullptr;
    CHECK(JS::ubi::Node(lazyScript).coarseType() == JS::ubi::CoarseType::Script);

    js::jit::JitCode* jitCode = nullptr;
    CHECK(JS::ubi::Node(jitCode).coarseType() == JS::ubi::CoarseType::Script);

    JSString* str = nullptr;
    CHECK(JS::ubi::Node(str).coarseType() == JS::ubi::CoarseType::String);

    // Test that the default when coarseType() is not overridden is Other.

    JS::Symbol* sym = nullptr;
    CHECK(JS::ubi::Node(sym).coarseType() == JS::ubi::CoarseType::Other);

    return true;
}
END_TEST(test_ubiCoarseType)

struct ExpectedEdge
{
    char from;
    char to;

    ExpectedEdge(FakeNode& fromNode, FakeNode& toNode)
        : from(fromNode.name)
        , to(toNode.name)
    { }
};

namespace js {

template <>
struct DefaultHasher<ExpectedEdge>
{
    using Lookup = ExpectedEdge;

    static HashNumber hash(const Lookup& l) {
        return mozilla::AddToHash(l.from, l.to);
    }

    static bool match(const ExpectedEdge& k, const Lookup& l) {
        return k.from == l.from && k.to == l.to;
    }
};

} // namespace js

BEGIN_TEST(test_ubiPostOrder)
{
    // Construct the following graph:
    //
    //                          .-----.
    //                          |     |
    //                  .-------|  r  |---------------.
    //                  |       |     |               |
    //                  |       '-----'               |
    //                  |                             |
    //               .--V--.                       .--V--.
    //               |     |                       |     |
    //        .------|  a  |------.           .----|  e  |----.
    //        |      |     |      |           |    |     |    |
    //        |      '--^--'      |           |    '-----'    |
    //        |         |         |           |               |
    //     .--V--.      |      .--V--.     .--V--.         .--V--.
    //     |     |      |      |     |     |     |         |     |
    //     |  b  |      '------|  c  |----->  f  |--------->  g  |
    //     |     |             |     |     |     |         |     |
    //     '-----'             '-----'     '-----'         '-----'
    //        |                   |
    //        |      .-----.      |
    //        |      |     |      |
    //        '------>  d  <------'
    //               |     |
    //               '-----'
    //

    FakeNode r('r');
    FakeNode a('a');
    FakeNode b('b');
    FakeNode c('c');
    FakeNode d('d');
    FakeNode e('e');
    FakeNode f('f');
    FakeNode g('g');

    js::HashSet<ExpectedEdge> expectedEdges(cx);
    CHECK(expectedEdges.init());

    auto declareEdge = [&](FakeNode& from, FakeNode& to) {
        return from.addEdgeTo(to) && expectedEdges.putNew(ExpectedEdge(from, to));
    };

    CHECK(declareEdge(r, a));
    CHECK(declareEdge(r, e));
    CHECK(declareEdge(a, b));
    CHECK(declareEdge(a, c));
    CHECK(declareEdge(b, d));
    CHECK(declareEdge(c, a));
    CHECK(declareEdge(c, d));
    CHECK(declareEdge(c, f));
    CHECK(declareEdge(e, f));
    CHECK(declareEdge(e, g));
    CHECK(declareEdge(f, g));

    js::Vector<char, 8, js::SystemAllocPolicy> visited;
    {
        // Do a PostOrder traversal, starting from r. Accumulate the names of
        // the nodes we visit in `visited`. Remove edges we traverse from
        // `expectedEdges` as we find them to ensure that we only find each edge
        // once.

        JS::AutoCheckCannotGC nogc(rt);
        JS::ubi::PostOrder traversal(rt, nogc);
        CHECK(traversal.init());
        CHECK(traversal.addStart(&r));

        auto onNode = [&](const JS::ubi::Node& node) {
            return visited.append(node.as<FakeNode>()->name);
        };

        auto onEdge = [&](const JS::ubi::Node& origin, const JS::ubi::Edge& edge) {
            ExpectedEdge e(*origin.as<FakeNode>(), *edge.referent.as<FakeNode>());
            if (!expectedEdges.has(e)) {
                fprintf(stderr,
                        "Error: Unexpected edge from %c to %c!\n",
                        origin.as<FakeNode>()->name,
                        edge.referent.as<FakeNode>()->name);
                return false;
            }

            expectedEdges.remove(e);
            return true;
        };

        CHECK(traversal.traverse(onNode, onEdge));
    }

    fprintf(stderr, "visited.length() = %lu\n", (unsigned long) visited.length());
    for (size_t i = 0; i < visited.length(); i++)
        fprintf(stderr, "visited[%lu] = '%c'\n", (unsigned long) i, visited[i]);

    CHECK(visited.length() == 8);
    CHECK(visited[0] == 'g');
    CHECK(visited[1] == 'f');
    CHECK(visited[2] == 'e');
    CHECK(visited[3] == 'd');
    CHECK(visited[4] == 'c');
    CHECK(visited[5] == 'b');
    CHECK(visited[6] == 'a');
    CHECK(visited[7] == 'r');

    // We found all the edges we expected.
    CHECK(expectedEdges.count() == 0);

    return true;
}
END_TEST(test_ubiPostOrder)

BEGIN_TEST(test_JS_ubi_DominatorTree)
{
    // Construct the following graph:
    //
    //                                  .-----.
    //                                  |     <--------------------------------.
    //          .--------+--------------|  r  |--------------.                 |
    //          |        |              |     |              |                 |
    //          |        |              '-----'              |                 |
    //          |     .--V--.                             .--V--.              |
    //          |     |     |                             |     |              |
    //          |     |  b  |                             |  c  |--------.     |
    //          |     |     |                             |     |        |     |
    //          |     '-----'                             '-----'        |     |
    //       .--V--.     |                                   |        .--V--.  |
    //       |     |     |                                   |        |     |  |
    //       |  a  <-----+                                   |   .----|  g  |  |
    //       |     |     |                              .----'   |    |     |  |
    //       '-----'     |                              |        |    '-----'  |
    //          |        |                              |        |       |     |
    //       .--V--.     |    .-----.                .--V--.     |       |     |
    //       |     |     |    |     |                |     |     |       |     |
    //       |  d  <-----+---->  e  <----.           |  f  |     |       |     |
    //       |     |          |     |    |           |     |     |       |     |
    //       '-----'          '-----'    |           '-----'     |       |     |
    //          |     .-----.    |       |              |        |    .--V--.  |
    //          |     |     |    |       |              |      .-'    |     |  |
    //          '----->  l  |    |       |              |      |      |  j  |  |
    //                |     |    '--.    |              |      |      |     |  |
    //                '-----'       |    |              |      |      '-----'  |
    //                   |       .--V--. |              |   .--V--.      |     |
    //                   |       |     | |              |   |     |      |     |
    //                   '------->  h  |-'              '--->  i  <------'     |
    //                           |     |          .--------->     |            |
    //                           '-----'          |         '-----'            |
    //                              |          .-----.         |               |
    //                              |          |     |         |               |
    //                              '---------->  k  <---------'               |
    //                                         |     |                         |
    //                                         '-----'                         |
    //                                            |                            |
    //                                            '----------------------------'
    //
    // This graph has the following dominator tree:
    //
    //     r
    //     |-- a
    //     |-- b
    //     |-- c
    //     |   |-- f
    //     |   `-- g
    //     |       `-- j
    //     |-- d
    //     |   `-- l
    //     |-- e
    //     |-- i
    //     |-- k
    //     `-- h
    //
    // This graph and dominator tree are taken from figures 1 and 2 of "A Fast
    // Algorithm for Finding Dominators in a Flowgraph" by Lengauer et al:
    // http://www.cs.princeton.edu/courses/archive/spr03/cs423/download/dominators.pdf.

    FakeNode r('r');
    FakeNode a('a');
    FakeNode b('b');
    FakeNode c('c');
    FakeNode d('d');
    FakeNode e('e');
    FakeNode f('f');
    FakeNode g('g');
    FakeNode h('h');
    FakeNode i('i');
    FakeNode j('j');
    FakeNode k('k');
    FakeNode l('l');

    CHECK(r.addEdgeTo(a));
    CHECK(r.addEdgeTo(b));
    CHECK(r.addEdgeTo(c));
    CHECK(a.addEdgeTo(d));
    CHECK(b.addEdgeTo(a));
    CHECK(b.addEdgeTo(d));
    CHECK(b.addEdgeTo(e));
    CHECK(c.addEdgeTo(f));
    CHECK(c.addEdgeTo(g));
    CHECK(d.addEdgeTo(l));
    CHECK(e.addEdgeTo(h));
    CHECK(f.addEdgeTo(i));
    CHECK(g.addEdgeTo(i));
    CHECK(g.addEdgeTo(j));
    CHECK(h.addEdgeTo(e));
    CHECK(h.addEdgeTo(k));
    CHECK(i.addEdgeTo(k));
    CHECK(j.addEdgeTo(i));
    CHECK(k.addEdgeTo(r));
    CHECK(k.addEdgeTo(i));
    CHECK(l.addEdgeTo(h));

    mozilla::Maybe<JS::ubi::DominatorTree> maybeTree;
    {
        JS::AutoCheckCannotGC noGC(rt);
        maybeTree = JS::ubi::DominatorTree::Create(rt, noGC, &r);
    }

    CHECK(maybeTree.isSome());
    auto& tree = *maybeTree;

    // We return the null JS::ubi::Node for nodes that were not reachable in the
    // graph when computing the dominator tree.
    FakeNode m('m');
    CHECK(tree.getImmediateDominator(&m) == JS::ubi::Node());
    CHECK(tree.getDominatedSet(&m).isNothing());

    struct {
        FakeNode& dominated;
        FakeNode& dominator;
    } domination[] = {
        {r, r},
        {a, r},
        {b, r},
        {c, r},
        {d, r},
        {e, r},
        {f, c},
        {g, c},
        {h, r},
        {i, r},
        {j, g},
        {k, r},
        {l, d}
    };

    for (auto& relation : domination) {
        // Test immediate dominator.
        fprintf(stderr,
                "%c's immediate dominator is %c\n",
                relation.dominated.name,
                tree.getImmediateDominator(&relation.dominator).as<FakeNode>()->name);
        CHECK(tree.getImmediateDominator(&relation.dominated) == JS::ubi::Node(&relation.dominator));

        // Test the dominated set. Build up the expected dominated set based on
        // the set of nodes immediately dominated by this one in `domination`,
        // then iterate over the actual dominated set and check against the
        // expected set.

        auto& node = relation.dominated;
        fprintf(stderr, "Checking %c's dominated set:\n", node.name);

        js::HashSet<char> expectedDominatedSet(cx);
        CHECK(expectedDominatedSet.init());
        for (auto& rel : domination) {
            if (&rel.dominator == &node) {
                fprintf(stderr, "    Expecting %c\n", rel.dominated.name);
                CHECK(expectedDominatedSet.putNew(rel.dominated.name));
            }
        }

        auto maybeActualDominatedSet = tree.getDominatedSet(&node);
        CHECK(maybeActualDominatedSet.isSome());
        auto& actualDominatedSet = *maybeActualDominatedSet;

        for (const auto& dominated : actualDominatedSet) {
            fprintf(stderr, "    Found %c\n", dominated.as<FakeNode>()->name);
            CHECK(expectedDominatedSet.has(dominated.as<FakeNode>()->name));
            expectedDominatedSet.remove(dominated.as<FakeNode>()->name);
        }

        // Ensure we found them all and aren't still expecting nodes we never
        // got.
        CHECK(expectedDominatedSet.count() == 0);

        fprintf(stderr, "Done checking %c's dominated set.\n\n", node.name);
    }

    struct {
        FakeNode& node;
        JS::ubi::Node::Size retainedSize;
    } sizes[] = {
        {r, 13},
        {a, 1},
        {b, 1},
        {c, 4},
        {d, 2},
        {e, 1},
        {f, 1},
        {g, 2},
        {h, 1},
        {i, 1},
        {j, 1},
        {k, 1},
        {l, 1},
    };

    for (auto& expected : sizes) {
        JS::ubi::Node::Size actual = 0;
        CHECK(tree.getRetainedSize(&expected.node, nullptr, actual));
        CHECK(actual == expected.retainedSize);
    }

    return true;
}
END_TEST(test_JS_ubi_DominatorTree)

BEGIN_TEST(test_JS_ubi_Node_scriptFilename)
{
    JS::RootedValue val(cx);
    CHECK(evaluate("(function one() {                      \n"  // 1
                   "  return (function two() {             \n"  // 2
                   "    return (function three() {         \n"  // 3
                   "      return function four() {};       \n"  // 4
                   "    }());                              \n"  // 5
                   "  }());                                \n"  // 6
                   "}());                                  \n", // 7
                   "my-cool-filename.js",
                   1,
                   &val));

    CHECK(val.isObject());
    JS::RootedObject obj(cx, &val.toObject());

    CHECK(obj->is<JSFunction>());
    JS::RootedFunction func(cx, &obj->as<JSFunction>());

    JS::RootedScript script(cx, func->getOrCreateScript(cx));
    CHECK(script);
    CHECK(script->filename());

    JS::ubi::Node node(script);
    const char* filename = node.scriptFilename();
    CHECK(filename);
    CHECK(strcmp(filename, script->filename()) == 0);
    CHECK(strcmp(filename, "my-cool-filename.js") == 0);

    return true;
}
END_TEST(test_JS_ubi_Node_scriptFilename)