Add key lookup support for tables in Go (#7644)

* Add support for key lookup for tables in Go

* Run clang format

* Run go fmt on tests

* Remove TODO in tests

* Update LookupByKey API

* Update LookupByKey API

* Don't use resolvePointer in expectEq

* Use generated getters instead of reading values directly from buffer

* Fix typo

Co-authored-by: Derek Bailey <derekbailey@google.com>
This commit is contained in:
Michael Le
2022-11-22 14:08:19 -08:00
committed by GitHub
parent 1cba8b2b49
commit 60975d6f7e
10 changed files with 365 additions and 7 deletions

View File

@@ -21,6 +21,7 @@
#include <sstream>
#include <string>
#include "flatbuffers/base.h"
#include "flatbuffers/code_generators.h"
#include "flatbuffers/flatbuffers.h"
#include "flatbuffers/idl.h"
@@ -104,6 +105,7 @@ class GoGenerator : public BaseGenerator {
++it) {
tracked_imported_namespaces_.clear();
needs_math_import_ = false;
needs_bytes_import_ = false;
needs_imports = false;
std::string enumcode;
GenEnum(**it, &enumcode);
@@ -124,6 +126,7 @@ class GoGenerator : public BaseGenerator {
it != parser_.structs_.vec.end(); ++it) {
tracked_imported_namespaces_.clear();
needs_math_import_ = false;
needs_bytes_import_ = false;
std::string declcode;
GenStruct(**it, &declcode);
if (parser_.opts.one_file) {
@@ -158,6 +161,7 @@ class GoGenerator : public BaseGenerator {
};
std::set<const Definition *, NamespacePtrLess> tracked_imported_namespaces_;
bool needs_math_import_ = false;
bool needs_bytes_import_ = false;
// Most field accessors need to retrieve and test the field offset first,
// this is the prefix code for that.
@@ -489,6 +493,34 @@ class GoGenerator : public BaseGenerator {
code += "}\n\n";
}
void GetMemberOfVectorOfStructByKey(const StructDef &struct_def,
const FieldDef &field,
std::string *code_ptr) {
std::string &code = *code_ptr;
auto vectortype = field.value.type.VectorType();
FLATBUFFERS_ASSERT(vectortype.struct_def->has_key);
auto &vector_struct_fields = vectortype.struct_def->fields.vec;
auto kit =
std::find_if(vector_struct_fields.begin(), vector_struct_fields.end(),
[&](FieldDef *field) { return field->key; });
auto &key_field = **kit;
FLATBUFFERS_ASSERT(key_field.key);
GenReceiver(struct_def, code_ptr);
code += " " + namer_.Field(field) + "ByKey";
code += "(obj *" + TypeName(field);
code += ", key " + NativeType(key_field.value.type) + ") bool" +
OffsetPrefix(field);
code += "\t\tx := rcv._tab.Vector(o)\n";
code += "\t\treturn ";
code += "obj.LookupByKey(key, x, rcv._tab.Bytes)\n";
code += "\t}\n";
code += "\treturn false\n";
code += "}\n\n";
}
// Get the value of a vector's non-struct member.
void GetMemberOfVectorOfNonStruct(const StructDef &struct_def,
const FieldDef &field,
@@ -690,6 +722,12 @@ class GoGenerator : public BaseGenerator {
auto vectortype = field.value.type.VectorType();
if (vectortype.base_type == BASE_TYPE_STRUCT) {
GetMemberOfVectorOfStruct(struct_def, field, code_ptr);
// TODO(michaeltle): Support querying fixed struct by key.
// Currently, we only support keyed tables.
if (!vectortype.struct_def->fixed &&
vectortype.struct_def->has_key) {
GetMemberOfVectorOfStructByKey(struct_def, field, code_ptr);
}
} else {
GetMemberOfVectorOfNonStruct(struct_def, field, code_ptr);
}
@@ -824,6 +862,12 @@ class GoGenerator : public BaseGenerator {
GenStructAccessor(struct_def, field, code_ptr);
GenStructMutator(struct_def, field, code_ptr);
// TODO(michaeltle): Support querying fixed struct by key. Currently,
// we only support keyed tables.
if (!struct_def.fixed && field.key) {
GenKeyCompare(struct_def, field, code_ptr);
GenLookupByKey(struct_def, field, code_ptr);
}
}
// Generate builders
@@ -836,6 +880,79 @@ class GoGenerator : public BaseGenerator {
}
}
void GenKeyCompare(const StructDef &struct_def, const FieldDef &field,
std::string *code_ptr) {
FLATBUFFERS_ASSERT(struct_def.has_key);
FLATBUFFERS_ASSERT(field.key);
std::string &code = *code_ptr;
code += "func " + namer_.Type(struct_def) + "KeyCompare(";
code += "o1, o2 flatbuffers.UOffsetT, buf []byte) bool {\n";
code += "\tobj1 := &" + namer_.Type(struct_def) + "{}\n";
code += "\tobj2 := &" + namer_.Type(struct_def) + "{}\n";
code += "\tobj1.Init(buf, flatbuffers.UOffsetT(len(buf)) - o1)\n";
code += "\tobj2.Init(buf, flatbuffers.UOffsetT(len(buf)) - o2)\n";
if (IsString(field.value.type)) {
code += "\treturn string(obj1." + namer_.Function(field.name) + "()) < ";
code += "string(obj2." + namer_.Function(field.name) + "())\n";
} else {
code += "\treturn obj1." + namer_.Function(field.name) + "() < ";
code += "obj2." + namer_.Function(field.name) + "()\n";
}
code += "}\n\n";
}
void GenLookupByKey(const StructDef &struct_def, const FieldDef &field,
std::string *code_ptr) {
FLATBUFFERS_ASSERT(struct_def.has_key);
FLATBUFFERS_ASSERT(field.key);
std::string &code = *code_ptr;
GenReceiver(struct_def, code_ptr);
code += " LookupByKey(";
code += "key " + NativeType(field.value.type) + ", ";
code += "vectorLocation flatbuffers.UOffsetT, ";
code += "buf []byte) bool {\n";
code += "\tspan := flatbuffers.GetUOffsetT(buf[vectorLocation - 4:])\n";
code += "\tstart := flatbuffers.UOffsetT(0)\n";
code += "\tfor span != 0 {\n";
code += "\t\tmiddle := span / 2\n";
code += "\t\ttableOffset := flatbuffers.GetIndirectOffset(buf, ";
code += "vectorLocation+ 4 * (start + middle))\n";
code += "\t\tobj := &" + namer_.Type(struct_def) + "{}\n";
code += "\t\tobj.Init(buf, tableOffset)\n";
if (IsString(field.value.type)) {
code += "\t\tbKey := []byte(key)\n";
needs_bytes_import_ = true;
code +=
"\t\tcomp := bytes.Compare(obj." + namer_.Function(field.name) + "()";
code += ", bKey)\n";
} else {
code += "\t\tval := obj." + namer_.Function(field.name) + "()\n";
code += "\t\tcomp := 0\n";
code += "\t\tif val > key {\n";
code += "\t\t\tcomp = 1\n";
code += "\t\t} else if val < key {\n";
code += "\t\t\tcomp = -1\n";
code += "\t\t}\n";
}
code += "\t\tif comp > 0 {\n";
code += "\t\t\tspan = middle\n";
code += "\t\t} else if comp < 0 {\n";
code += "\t\t\tmiddle += 1\n";
code += "\t\t\tstart += middle\n";
code += "\t\t\tspan -= middle\n";
code += "\t\t} else {\n";
code += "\t\t\trcv.Init(buf, tableOffset)\n";
code += "\t\t\treturn true\n";
code += "\t\t}\n";
code += "\t}\n";
code += "\treturn false\n";
code += "}\n\n";
}
void GenNativeStruct(const StructDef &struct_def, std::string *code_ptr) {
std::string &code = *code_ptr;
@@ -1354,9 +1471,10 @@ class GoGenerator : public BaseGenerator {
code += "package " + name_space_name + "\n\n";
if (needs_imports) {
code += "import (\n";
if (is_enum) { code += "\t\"strconv\"\n\n"; }
if (needs_bytes_import_) code += "\t\"bytes\"\n";
// math is needed to support non-finite scalar default values.
if (needs_math_import_) { code += "\t\"math\"\n\n"; }
if (needs_math_import_) { code += "\t\"math\"\n"; }
if (is_enum) { code += "\t\"strconv\"\n"; }
if (!parser_.opts.go_import.empty()) {
code += "\tflatbuffers \"" + parser_.opts.go_import + "\"\n";
} else {