class for referring to definitions of variables or record fields

src/CMakeLists.txt

@@ -322,6 +322,7 @@ set(MAIN_SRCS
     plugin/Manager.cc
     plugin/Plugin.cc
 
+    script_opt/DefItem.cc
     script_opt/Expr.cc
     script_opt/Inline.cc
     script_opt/ProfileFunc.cc

src/script_opt/DefItem.cc

@@ -0,0 +1,154 @@
+// See the file "COPYING" in the main distribution directory for copyright.
+
+#include "DefItem.h"
+#include "Expr.h"
+
+
+namespace zeek::detail {
+
+
+DefinitionItem::DefinitionItem(const ID* _id)
+: name(_id->Name())
+	{
+	is_id = true;
+	id = _id;
+	di = nullptr;
+	field_name = nullptr;
+
+	t = id->GetType();
+
+	CheckForRecord();
+	}
+
+DefinitionItem::DefinitionItem(const DefinitionItem* _di,
+				const char* _field_name, TypePtr _t)
+	{
+	is_id = false;
+	id = nullptr;
+	di = _di;
+	field_name = _field_name;
+
+	t = std::move(_t);
+
+	name += di->Name();
+	name += '$';
+	name += field_name;
+
+	CheckForRecord();
+	}
+
+std::shared_ptr<DefinitionItem>
+DefinitionItem::FindField(const char* field) const
+	{
+	if ( ! IsRecord() )
+		return nullptr;
+
+	auto offset = rt->FieldOffset(field);
+
+	return FindField(offset);
+	}
+
+std::shared_ptr<DefinitionItem> DefinitionItem::FindField(int offset) const
+	{
+	if ( ! IsRecord() )
+		return nullptr;
+
+	return (*fields)[offset];
+	}
+
+std::shared_ptr<DefinitionItem>
+DefinitionItem::CreateField(const char* field, TypePtr t)
+	{
+	auto offset = rt->FieldOffset(field);
+
+	if ( (*fields)[offset] )
+		return (*fields)[offset];
+
+	(*fields)[offset] =
+		std::make_shared<DefinitionItem>(this, field, std::move(t));
+
+	return (*fields)[offset];
+	}
+
+std::shared_ptr<DefinitionItem>
+DefinitionItem::CreateField(int offset, TypePtr t)
+	{
+	if ( (*fields)[offset] )
+		return (*fields)[offset];
+
+	auto field = rt->FieldName(offset);
+
+	(*fields)[offset] =
+		std::make_shared<DefinitionItem>(this, field, std::move(t));
+
+	return (*fields)[offset];
+	}
+
+void DefinitionItem::CheckForRecord()
+	{
+	if ( ! IsRecord() )
+		{
+		rt = nullptr;
+		return;
+		}
+
+	rt = t->AsRecordType();
+	num_fields = rt->NumFields();
+	fields = std::vector<std::shared_ptr<DefinitionItem>>(num_fields);
+	}
+
+
+std::shared_ptr<DefinitionItem> DefItemMap::GetExprDI(const Expr* expr)
+	{
+	if ( expr->Tag() == EXPR_NAME )
+		{
+		auto id_e = expr->AsNameExpr();
+		auto id = id_e->Id();
+		return GetID_DI(id);
+		}
+
+	else if ( expr->Tag() == EXPR_FIELD )
+		{
+		auto f = expr->AsFieldExpr();
+		auto r = f->Op();
+
+		auto r_def = GetExprDI(r);
+
+		if ( ! r_def )
+			return nullptr;
+
+		auto field = f->FieldName();
+		return r_def->FindField(field);
+		}
+
+	else
+		return nullptr;
+	}
+
+std::shared_ptr<DefinitionItem> DefItemMap::GetID_DI(const ID* id)
+	{
+	auto di = i2d.find(id);
+	if ( di == i2d.end() )
+		{
+		auto new_entry = std::make_shared<DefinitionItem>(id);
+		i2d[id] = new_entry;
+		return new_entry;
+		}
+	else
+		return di->second;
+	}
+
+const DefinitionItem* DefItemMap::GetConstID_DI(const ID* id) const
+	{
+	auto di = i2d.find(id);
+	return di == i2d.end() ? nullptr : di->second.get();
+	}
+
+const DefinitionItem* DefItemMap::GetConstID_DI(const DefinitionItem* di,
+						const char* field_name) const
+	{
+	return di->FindField(field_name).get();
+	}
+
+
+} // zeek::detail

src/script_opt/DefItem.h

@@ -0,0 +1,107 @@
+// See the file "COPYING" in the main distribution directory for copyright.
+
+#pragma once
+
+#include "zeek/script_opt/DefPoint.h"
+#include "zeek/ID.h"
+#include "zeek/Type.h"
+
+
+namespace zeek::detail {
+
+
+// A definition item is a Zeek script entity that can be assigned to.
+// Currently, we track variables and record fields; the latter can
+// be nested (for example, a field that's in a record that itself is
+// a field in another record).  In principle we could try to track
+// table or vector elements, but that's only going to be feasible for
+// constant indices, so presumably not much bang-for-the-buck.
+//
+// For script optimization, we only need to track variables, and we could
+// considerably simplify the code by doing so.  However, there's long
+// been a desire to be able to statically determine that a record field
+// will be used without first having been set, hence we go the more
+// complicated route here.
+
+class DefinitionItem {
+public:
+	// Constructor for the simple case of tracking assignments to
+	// a variable.
+	DefinitionItem(const ID* _id);
+
+	// The more complicated case of assigning to a field in a record
+	// (which itself might be a field in a record).
+	DefinitionItem(const DefinitionItem* _di, const char* _field_name,
+			TypePtr _t);
+
+	const char* Name() const	{ return name.c_str(); }
+
+	TypePtr GetType() const		{ return t; }
+	bool IsRecord() const		{ return t->Tag() == TYPE_RECORD; }
+
+	// The identifier to which this item ultimately belongs.
+	const ID* RootID() const	{ return di ? di->RootID() : id; }
+
+	// For this definition item, look for a field corresponding
+	// to the given name or offset.  Nil if the field has not (yet)
+	// been created.
+	std::shared_ptr<DefinitionItem> FindField(const char* field) const;
+	std::shared_ptr<DefinitionItem> FindField(int offset) const;
+
+	// Start tracking a field in this definition item with the
+	// given name or offset, returning the associated item.
+	//
+	// If the field already exists, then it's simply returned.
+	std::shared_ptr<DefinitionItem> CreateField(const char* field,
+							TypePtr t);
+	std::shared_ptr<DefinitionItem> CreateField(int offset, TypePtr t);
+
+protected:
+	void CheckForRecord();
+
+	bool is_id;
+	const ID* id;
+	const DefinitionItem* di;
+	const char* field_name;
+
+	TypePtr t;
+	std::string name;
+
+	const RecordType* rt;
+
+	// If present, tracks definition items for a record's fields as
+	// these are seen (i.e., as they are entered via CreateField()).
+	std::optional<std::vector<std::shared_ptr<DefinitionItem>>> fields;
+	int num_fields;
+};
+
+// For a given identifier, locates its associated definition item.
+typedef std::unordered_map<const ID*, std::shared_ptr<DefinitionItem>>
+		ID_to_DI_Map;
+
+// Class for managing a set of IDs and their associated definition items.
+class DefItemMap {
+public:
+	// Gets the definition for either a name or a record field reference.
+	// Returns nil if "expr" lacks such a form, or if there isn't
+	// any such definition.
+	std::shared_ptr<DefinitionItem> GetExprDI(const Expr* expr);
+
+	// Returns the definition item for a given ID; creates it if
+	// it doesn't already exist.
+	std::shared_ptr<DefinitionItem> GetID_DI(const ID* id);
+
+	// Returns the definition item for a given ID, or nil if it
+	// doesn't exist.
+	const DefinitionItem* GetConstID_DI(const ID* id) const;
+
+	// The same for a record field for a given definition item.
+	const DefinitionItem* GetConstID_DI(const DefinitionItem* di,
+						const char* field_name) const;
+
+protected:
+	ID_to_DI_Map i2d;
+};
+
+
+} // zeek::detail