luau/tests/AstQueryDsl.h
vegorov-rbx 76070f8da2
Sync to upstream/release/549 (#707)
* Reoptimized math.min/max/bit32 builtins assuming at least 2 arguments are used (1-2% lift on some benchmarks)
* Type errors that mention function types no longer have redundant parenthesis around return type
* Luau REPL now supports --compile=remarks which displays the source code with optimization remarks embedded as comments
* Builtin calls are slightly faster when called with 1-2 arguments (~1% improvement in some benchmarks)
2022-10-14 12:48:41 -07:00

83 lines
2 KiB
C++

// This file is part of the Luau programming language and is licensed under MIT License; see LICENSE.txt for details
#pragma once
#include "Luau/Ast.h"
#include "Luau/Common.h"
#include <vector>
#include <type_traits>
namespace Luau
{
struct Nth
{
int classIndex;
int nth;
};
template<typename T>
Nth nth(int nth = 1)
{
static_assert(std::is_base_of_v<AstNode, T>, "T must be a derived class of AstNode");
LUAU_ASSERT(nth > 0); // Did you mean to use `nth<T>(1)`?
return Nth{T::ClassIndex(), nth};
}
struct FindNthOccurenceOf : public AstVisitor
{
Nth requestedNth;
int currentOccurrence = 0;
AstNode* theNode = nullptr;
FindNthOccurenceOf(Nth nth);
bool checkIt(AstNode* n);
bool visit(AstNode* n) override;
bool visit(AstType* n) override;
bool visit(AstTypePack* n) override;
};
/** DSL querying of the AST.
*
* Given an AST, one can query for a particular node directly without having to manually unwrap the tree, for example:
*
* ```
* if a and b then
* print(a + b)
* end
*
* function f(x, y)
* return x + y
* end
* ```
*
* There are numerous ways to access the second AstExprBinary.
* 1. Luau::query<AstExprBinary>(block, {nth<AstStatFunction>(), nth<AstExprBinary>()})
* 2. Luau::query<AstExprBinary>(Luau::query<AstStatFunction>(block))
* 3. Luau::query<AstExprBinary>(block, {nth<AstExprBinary>(2)})
*/
template<typename T, int N = 1>
T* query(AstNode* node, const std::vector<Nth>& nths = {nth<T>(N)})
{
static_assert(std::is_base_of_v<AstNode, T>, "T must be a derived class of AstNode");
// If a nested query call fails to find the node in question, subsequent calls can propagate rather than trying to do more.
// This supports `query(query(...))`
for (Nth nth : nths)
{
if (!node)
return nullptr;
FindNthOccurenceOf finder{nth};
node->visit(&finder);
node = finder.theNode;
}
return node ? node->as<T>() : nullptr;
}
} // namespace Luau