Tag notifications & tag type list optimized

9 files changed, 112 insertions, 13 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 0a775ac8..1728774d 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -908,6 +908,8 @@ namespace BinaryNinja
 	struct DataVariable;
 	class Symbol;
 	class Tag;
+	class TagType;
+	struct TagReference;
 
 	class BinaryDataNotification
 	{
@@ -928,6 +930,10 @@ namespace BinaryNinja
 		static void SymbolUpdatedCallback(void* ctxt, BNBinaryView* view, BNSymbol* sym);
 		static void SymbolRemovedCallback(void* ctxt, BNBinaryView* view, BNSymbol* sym);
 		static void DataMetadataUpdatedCallback(void* ctxt, BNBinaryView* object, uint64_t offset);
+		static void TagTypeUpdatedCallback(void* ctxt, BNBinaryView* object, BNTagType* tagType);
+		static void TagAddedCallback(void* ctxt, BNBinaryView* object, BNTagReference* tagRef);
+		static void TagUpdatedCallback(void* ctxt, BNBinaryView* object, BNTagReference* tagRef);
+		static void TagRemovedCallback(void* ctxt, BNBinaryView* object, BNTagReference* tagRef);
 		static void StringFoundCallback(void* ctxt, BNBinaryView* data, BNStringType type, uint64_t offset, size_t len);
 		static void StringRemovedCallback(void* ctxt, BNBinaryView* data, BNStringType type, uint64_t offset, size_t len);
 		static void TypeDefinedCallback(void* ctxt, BNBinaryView* data, BNQualifiedName* name, BNType* type);
@@ -950,6 +956,10 @@ namespace BinaryNinja
 		virtual void OnDataVariableRemoved(BinaryView* view, const DataVariable& var) { (void)view; (void)var; }
 		virtual void OnDataVariableUpdated(BinaryView* view, const DataVariable& var) { (void)view; (void)var; }
 		virtual void OnDataMetadataUpdated(BinaryView* view, uint64_t offset) { (void)view; (void)offset; }
+		virtual void OnTagTypeUpdated(BinaryView* view, Ref<TagType> tagTypeRef) { (void)view; (void)tagTypeRef; }
+		virtual void OnTagAdded(BinaryView* view, const TagReference& tagRef) { (void)view; (void)tagRef; }
+		virtual void OnTagUpdated(BinaryView* view, const TagReference& tagRef) { (void)view; (void)tagRef; }
+		virtual void OnTagRemoved(BinaryView* view, const TagReference& tagRef) { (void)view; (void)tagRef; }
 		virtual void OnSymbolAdded(BinaryView* view, Symbol* sym) { (void)view; (void)sym; }
 		virtual void OnSymbolUpdated(BinaryView* view, Symbol* sym) { (void)view; (void)sym; }
 		virtual void OnSymbolRemoved(BinaryView* view, Symbol* sym) { (void)view; (void)sym; }
@@ -1559,8 +1569,8 @@ namespace BinaryNinja
 		Ref<TagType> GetTagType(const std::string& name, TagType::Type type);
 		std::vector<Ref<TagType>> GetTagTypes();
 
-		void AddTag(Ref<Tag> tag);
-		void RemoveTag(Ref<Tag> tag);
+		void AddTag(Ref<Tag> tag, bool user = false);
+		void RemoveTag(Ref<Tag> tag, bool user = false);
 		Ref<Tag> GetTag(uint64_t tagId);
 
 		std::vector<TagReference> GetAllTagReferences();
@@ -1569,6 +1579,9 @@ namespace BinaryNinja
 		std::vector<TagReference> GetAllTagReferencesOfType(Ref<TagType> tagType);
 		std::vector<TagReference> GetTagReferencesOfType(Ref<TagType> tagType);
 
+		size_t GetTagReferencesOfTypeCount(Ref<TagType> tagType);
+		size_t GetAllTagReferencesOfTypeCount(Ref<TagType> tagType);
+
 		std::vector<TagReference> GetDataTagReferences();
 		std::vector<Ref<Tag>> GetDataTags(uint64_t addr);
 		std::vector<Ref<Tag>> GetDataTagsOfType(uint64_t addr, Ref<TagType> tagType);
@@ -2880,6 +2893,7 @@ namespace BinaryNinja
 
 		std::vector<TagReference> GetAllTagReferences();
 		std::vector<TagReference> GetTagReferencesOfType(Ref<TagType> tagType);
+		size_t GetTagReferencesOfTypeCount(Ref<TagType> tagType);
 
 		std::vector<TagReference> GetAddressTagReferences();
 		std::vector<Ref<Tag>> GetAddressTags(Architecture* arch, uint64_t addr);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 1e0ba256..ada6c81e 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -139,6 +139,7 @@ extern "C"
 	struct BNStructure;
 	struct BNTagType;
 	struct BNTag;
+	struct BNTagReference;
 	struct BNNamedTypeReference;
 	struct BNEnumeration;
 	struct BNCallingConvention;
@@ -1080,6 +1081,10 @@ extern "C"
 		void (*dataVariableRemoved)(void* ctxt, BNBinaryView* view, BNDataVariable* var);
 		void (*dataVariableUpdated)(void* ctxt, BNBinaryView* view, BNDataVariable* var);
 		void (*dataMetadataUpdated)(void* ctxt, BNBinaryView* view, uint64_t offset);
+		void (*tagTypeUpdated)(void* ctxt, BNBinaryView* view, BNTagType* tagType);
+		void (*tagAdded)(void* ctxt, BNBinaryView* view, BNTagReference* tagRef);
+		void (*tagUpdated)(void* ctxt, BNBinaryView* view, BNTagReference* tagRef);
+		void (*tagRemoved)(void* ctxt, BNBinaryView* view, BNTagReference* tagRef);
 		void (*symbolAdded)(void* ctxt, BNBinaryView* view, BNSymbol* sym);
 		void (*symbolUpdated)(void* ctxt, BNBinaryView* view, BNSymbol* sym);
 		void (*symbolRemoved)(void* ctxt, BNBinaryView* view, BNSymbol* sym);
@@ -2964,9 +2969,9 @@ extern "C"
 	BINARYNINJACOREAPI BNTagType* BNGetTagTypeWithType(BNBinaryView* view, const char* name, BNTagTypeType type);
 	BINARYNINJACOREAPI BNTagType** BNGetTagTypes(BNBinaryView* view, size_t* count);
 
-	BINARYNINJACOREAPI void BNAddTag(BNBinaryView* view, BNTag* tag);
+	BINARYNINJACOREAPI void BNAddTag(BNBinaryView* view, BNTag* tag, bool user);
 	BINARYNINJACOREAPI BNTag* BNGetTag(BNBinaryView* view, uint64_t tagId);
-	BINARYNINJACOREAPI void BNRemoveTag(BNBinaryView* view, BNTag* tag);
+	BINARYNINJACOREAPI void BNRemoveTag(BNBinaryView* view, BNTag* tag, bool user);
 
 	BINARYNINJACOREAPI BNTagReference* BNGetAllTagReferences(BNBinaryView* view, size_t* count);
 	BINARYNINJACOREAPI BNTagReference* BNGetAllAddressTagReferences(BNBinaryView* view, size_t* count);
@@ -2984,6 +2989,9 @@ extern "C"
 	BINARYNINJACOREAPI void BNRemoveUserDataTag(BNBinaryView* view, uint64_t addr, BNTag* tag);
 	BINARYNINJACOREAPI void BNRemoveTagReference(BNBinaryView* view, BNTagReference ref);
 
+	BINARYNINJACOREAPI size_t BNGetTagReferencesOfTypeCount(BNBinaryView* view, BNTagType* tagType);
+	BINARYNINJACOREAPI size_t BNGetAllTagReferencesOfTypeCount(BNBinaryView* view, BNTagType* tagType);
+
 	BINARYNINJACOREAPI BNTagReference* BNGetFunctionAllTagReferences(BNFunction* func, size_t* count);
 	BINARYNINJACOREAPI BNTagReference* BNGetFunctionTagReferencesOfType(BNFunction* func, BNTagType* tagType, size_t* count);
 
diff --git a/binaryview.cpp b/binaryview.cpp
index dc2212a6..5cec7b10 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -122,6 +122,39 @@ void BinaryDataNotification::DataMetadataUpdatedCallback(void* ctxt, BNBinaryVie
 }
 
 
+void BinaryDataNotification::TagTypeUpdatedCallback(void* ctxt, BNBinaryView* object, BNTagType* tagType)
+{
+	BinaryDataNotification* notify = (BinaryDataNotification*)ctxt;
+	Ref<BinaryView> view = new BinaryView(BNNewViewReference(object));
+	Ref<TagType> tagTypeRef = new TagType(BNNewTagTypeReference(tagType));
+	notify->OnTagTypeUpdated(view, tagTypeRef);
+}
+
+
+void BinaryDataNotification::TagAddedCallback(void* ctxt, BNBinaryView* object, BNTagReference* tagRef)
+{
+	BinaryDataNotification* notify = (BinaryDataNotification*)ctxt;
+	Ref<BinaryView> view = new BinaryView(BNNewViewReference(object));
+	notify->OnTagAdded(view, TagReference(*tagRef));
+}
+
+
+void BinaryDataNotification::TagUpdatedCallback(void* ctxt, BNBinaryView* object, BNTagReference* tagRef)
+{
+	BinaryDataNotification* notify = (BinaryDataNotification*)ctxt;
+	Ref<BinaryView> view = new BinaryView(BNNewViewReference(object));
+	notify->OnTagUpdated(view, TagReference(*tagRef));
+}
+
+
+void BinaryDataNotification::TagRemovedCallback(void* ctxt, BNBinaryView* object, BNTagReference* tagRef)
+{
+	BinaryDataNotification* notify = (BinaryDataNotification*)ctxt;
+	Ref<BinaryView> view = new BinaryView(BNNewViewReference(object));
+	notify->OnTagRemoved(view, TagReference(*tagRef));
+}
+
+
 void BinaryDataNotification::SymbolAddedCallback(void* ctxt, BNBinaryView* object, BNSymbol* symobj)
 {
 	BinaryDataNotification* notify = (BinaryDataNotification*)ctxt;
@@ -197,6 +230,10 @@ BinaryDataNotification::BinaryDataNotification()
 	m_callbacks.dataVariableRemoved = DataVariableRemovedCallback;
 	m_callbacks.dataVariableUpdated = DataVariableUpdatedCallback;
 	m_callbacks.dataMetadataUpdated = DataMetadataUpdatedCallback;
+	m_callbacks.tagTypeUpdated = TagTypeUpdatedCallback;
+	m_callbacks.tagAdded = TagAddedCallback;
+	m_callbacks.tagUpdated = TagUpdatedCallback;
+	m_callbacks.tagRemoved = TagRemovedCallback;
 	m_callbacks.symbolAdded = SymbolAddedCallback;
 	m_callbacks.symbolUpdated = SymbolUpdatedCallback;
 	m_callbacks.symbolRemoved = SymbolRemovedCallback;
@@ -2047,15 +2084,15 @@ std::vector<Ref<TagType>> BinaryView::GetTagTypes()
 }
 
 
-void BinaryView::AddTag(Ref<Tag> tag)
+void BinaryView::AddTag(Ref<Tag> tag, bool user)
 {
-	BNAddTag(m_object, tag->GetObject());
+	BNAddTag(m_object, tag->GetObject(), user);
 }
 
 
-void BinaryView::RemoveTag(Ref<Tag> tag)
+void BinaryView::RemoveTag(Ref<Tag> tag, bool user)
 {
-	BNRemoveTag(m_object, tag->GetObject());
+	BNRemoveTag(m_object, tag->GetObject(), user);
 }
 
 
@@ -2109,6 +2146,18 @@ std::vector<TagReference> BinaryView::GetTagReferencesOfType(Ref<TagType> tagTyp
 }
 
 
+size_t BinaryView::GetAllTagReferencesOfTypeCount(Ref<TagType> tagType)
+{
+	return BNGetAllTagReferencesOfTypeCount(m_object, tagType->GetObject());
+}
+
+
+size_t BinaryView::GetTagReferencesOfTypeCount(Ref<TagType> tagType)
+{
+	return BNGetTagReferencesOfTypeCount(m_object, tagType->GetObject());
+}
+
+
 std::vector<TagReference> BinaryView::GetDataTagReferences()
 {
 	size_t count;
diff --git a/python/binaryview.py b/python/binaryview.py
index ee4a4fb3..4e7b46a7 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -3468,7 +3468,13 @@ class BinaryView(object):
 		core.BNFreeTagTypeList(types, count.value)
 		return result
 
-	def create_tag(self, type, data):
+	def create_user_tag(self, type, data):
+		return self.create_tag(type, data, True)
+
+	def create_auto_tag(self, type, data):
+		return self.create_tag(type, data, False)
+
+	def create_tag(self, type, data, user=True):
 		"""
 		``create_tag`` creates a new Tag object but does not add it anywhere
 
@@ -3483,7 +3489,7 @@ class BinaryView(object):
 			>>> bv.add_user_data_tag(here, tag)
 			>>>
 		"""
-		tag = Tag(core.BNCreateTag(type.handle, data))
+		tag = Tag(core.BNCreateTag(type.handle, data, user))
 		core.BNAddTag(self.handle, tag.handle)
 		return tag
 
diff --git a/ui/flowgraphwidget.h b/ui/flowgraphwidget.h
index 9c4dfcd0..0d1c6355 100644
--- a/ui/flowgraphwidget.h
+++ b/ui/flowgraphwidget.h
@@ -182,6 +182,7 @@ public:
 	virtual void OnAnalysisFunctionUpdated(BinaryNinja::BinaryView* data, BinaryNinja::Function* func) override;
 	virtual void OnAnalysisFunctionUpdateRequested(BinaryNinja::BinaryView* data, BinaryNinja::Function* func) override;
 	virtual void OnDataMetadataUpdated(BinaryNinja::BinaryView* data, uint64_t offset) override;
+	virtual void OnTagUpdated(BinaryNinja::BinaryView* data, const BinaryNinja::TagReference& tagRef) override;
 
 	void setInitialGraph(FlowGraphRef graph);
 	void setInitialGraph(FlowGraphRef graph, uint64_t addr);
diff --git a/ui/linearview.h b/ui/linearview.h
index 7a7f9ad8..adc14f2e 100644
--- a/ui/linearview.h
+++ b/ui/linearview.h
@@ -213,6 +213,7 @@ public:
 	virtual void OnDataVariableRemoved(BinaryNinja::BinaryView* view, const BinaryNinja::DataVariable& var) override;
 	virtual void OnDataVariableUpdated(BinaryNinja::BinaryView* view, const BinaryNinja::DataVariable& var) override;
 	virtual void OnDataMetadataUpdated(BinaryNinja::BinaryView* view, uint64_t offset) override;
+	virtual void OnTagUpdated(BinaryNinja::BinaryView* data, const BinaryNinja::TagReference& tagRef) override;
 
 	virtual void updateFonts() override;
 
diff --git a/ui/taglist.h b/ui/taglist.h
index c2a5bd8a..11e2333e 100644
--- a/ui/taglist.h
+++ b/ui/taglist.h
@@ -140,7 +140,13 @@ protected:
 	void selectFirstItem() override;
 	void activateFirstItem() override;
 
-	virtual void OnDataMetadataUpdated(BinaryNinja::BinaryView*, uint64_t) override;
+	virtual void OnTagAdded(BinaryNinja::BinaryView*, const BinaryNinja::TagReference&) override;
+	virtual void OnTagUpdated(BinaryNinja::BinaryView*, const BinaryNinja::TagReference&) override;
+	virtual void OnTagRemoved(BinaryNinja::BinaryView*, const BinaryNinja::TagReference&) override;
+	virtual void OnTagTypeUpdated(BinaryNinja::BinaryView*, TagTypeRef) override;
+
+	virtual void showEvent(QShowEvent *event) override;
+	virtual void hideEvent(QHideEvent *event) override;
 
 private Q_SLOTS:
 	void hoverTimerEvent();
diff --git a/ui/tagtypelist.h b/ui/tagtypelist.h
index cc7015a9..98ad2305 100644
--- a/ui/tagtypelist.h
+++ b/ui/tagtypelist.h
@@ -10,14 +10,16 @@
 #include "dockhandler.h"
 #include "viewframe.h"
 
+#define TAGS_UPDATE_CHECK_INTERVAL 200
 
-class BINARYNINJAUIAPI TagTypeListModel: public QAbstractItemModel
+class BINARYNINJAUIAPI TagTypeListModel: public QAbstractItemModel, public BinaryNinja::BinaryDataNotification
 {
 	Q_OBJECT
 
 	QWidget* m_owner;
 	BinaryViewRef m_data;
 	std::vector<TagTypeRef> m_refs;
+	std::map<std::string, uint64_t> m_count;
 
 public:
 	TagTypeListModel(QWidget* parent, BinaryViewRef data);
@@ -36,7 +38,13 @@ public:
 	virtual void sort(int column, Qt::SortOrder order) override;
 	virtual Qt::ItemFlags flags(const QModelIndex& i) const override;
 
+	virtual void OnTagAdded(BinaryNinja::BinaryView*, const BinaryNinja::TagReference&) override;
+	virtual void OnTagRemoved(BinaryNinja::BinaryView*, const BinaryNinja::TagReference&) override;
+
 	bool setModelData(const std::vector<TagTypeRef>& refs, QItemSelectionModel* selectionModel, int sortColumn, Qt::SortOrder sortOrder, bool& selectionUpdated);
+
+Q_SIGNALS:
+	void needRepaint();
 };
 
 
@@ -77,11 +85,15 @@ class BINARYNINJAUIAPI TagTypeList: public QTableView, public DockContextHandler
 
 	QTimer* m_updateTimer;
 	bool m_needsUpdate;
+	bool m_needsRepaint;
 
 protected:
 	virtual void contextMenuEvent(QContextMenuEvent* event) override;
 
-	virtual void OnDataMetadataUpdated(BinaryNinja::BinaryView*, uint64_t) override;
+	virtual void showEvent(QShowEvent *event) override;
+	virtual void hideEvent(QHideEvent *event) override;
+
+	virtual void OnTagTypeUpdated(BinaryNinja::BinaryView*, const TagTypeRef) override;
 
 private:
 	void createTagType();
@@ -89,6 +101,7 @@ private:
 
 private Q_SLOTS:
 	void updateTimerEvent();
+	void needRepaint();
 
 public:
 	TagTypeList(QWidget* parent, ViewFrame* view, BinaryViewRef data, Menu* menu = nullptr);
diff --git a/ui/tokenizedtextview.h b/ui/tokenizedtextview.h
index 95c47666..7ef15ae9 100644
--- a/ui/tokenizedtextview.h
+++ b/ui/tokenizedtextview.h
@@ -138,6 +138,7 @@ public:
 	virtual void OnDataVariableRemoved(BinaryNinja::BinaryView* view, const BinaryNinja::DataVariable& var) override;
 	virtual void OnDataVariableUpdated(BinaryNinja::BinaryView* view, const BinaryNinja::DataVariable& var) override;
 	virtual void OnDataMetadataUpdated(BinaryNinja::BinaryView* view, uint64_t offset) override;
+	virtual void OnTagUpdated(BinaryNinja::BinaryView* data, const BinaryNinja::TagReference& tagRef) override;
 
 	virtual void updateFonts() override;