Make item additions undo-able

4 files changed, 67 insertions, 17 deletions

diff --git a/src/commands.cpp b/src/commands.cpp
index 4930ac8..c6eb692 100644
--- a/src/commands.cpp
+++ b/src/commands.cpp
@@ -16,6 +16,7 @@
 
 #include <QDebug>
 #include "commands.h"
+#include "diagramdocument.h"
 #include "databasetable.h"
 #include "column.h"
 
@@ -143,3 +144,29 @@ MoveItemCommand::undo()
 {
 	m_item->setPos(m_oldPos);
 }
+
+
+AddItemCommand::AddItemCommand(DiagramDocument *document, DiagramItem *item, QUndoCommand *parent)
+	: QUndoCommand(parent), m_document(document), m_item(item), m_owner(true)
+{
+}
+
+AddItemCommand::~AddItemCommand()
+{
+	if (m_owner)
+		delete m_item;
+}
+
+void
+AddItemCommand::redo()
+{
+	m_document->addItem(m_item);
+	m_owner = false;
+}
+
+void
+AddItemCommand::undo()
+{
+	m_document->removeItem(m_item);
+	m_owner = true;
+}
diff --git a/src/commands.h b/src/commands.h
index 142b62c..617193c 100644
--- a/src/commands.h
+++ b/src/commands.h
@@ -98,4 +98,18 @@ private:
 	QPointF m_oldPos, m_newPos;
 };
 
+class AddItemCommand : public QUndoCommand
+{
+public:
+	AddItemCommand(DiagramDocument *document, DiagramItem *item, QUndoCommand *parent = 0);
+	~AddItemCommand();
+	void undo();
+	void redo();
+
+private:
+	DiagramDocument *m_document;
+	DiagramItem *m_item;
+	bool m_owner;
+};
+
 #endif
diff --git a/src/diagramdocument.cpp b/src/diagramdocument.cpp
index fe8c53f..d0a8b4c 100644
--- a/src/diagramdocument.cpp
+++ b/src/diagramdocument.cpp
@@ -52,6 +52,18 @@ DiagramDocument::updatePositions(DatabaseTable *table)
 	}
 }
 
+template <class T> QList<T *>
+DiagramDocument::itemsByType()
+{
+	QList<T *> result;
+	foreach(QGraphicsItem *item, items()) {
+		if (item->type() == T::Type) {
+			result.append(static_cast<T *>(item));
+		}
+	}
+	return result;
+}
+
 void
 DiagramDocument::itemMoved(DiagramItem *item)
 {
@@ -77,8 +89,7 @@ DiagramDocument::mousePressEvent(QGraphicsSceneMouseEvent *event)
 		DatabaseTable *table = new DatabaseTable();
 		table->setPos(event->scenePos());
 		table->setZValue(10.0);
-		m_tables << table;
-		addItem(table);
+		undoStack()->push(new AddItemCommand(this, table));
 		clearSelection();
 		table->setSelected(true);
 		setMode(Select);
@@ -121,7 +132,6 @@ DiagramDocument::mouseReleaseEvent(QGraphicsSceneMouseEvent *event)
 			relation->setTarget(target);
 			relation->setZValue(1.0);
 			addItem(relation);
-			m_relations << relation;
 		}
 
 		delete m_line;
@@ -157,10 +167,8 @@ DiagramDocument::deleteSelectedItems()
 		if (table) {
 			foreach (DatabaseRelationship *relation, findTableRelations(table)) {
 				removeItem(relation);
-				m_relations.removeAll(relation);
 			}
 			removeItem(table);
-			m_tables.removeAll(table);
 		}
 	}
 }
@@ -169,7 +177,7 @@ QList<DatabaseRelationship *>
 DiagramDocument::findTableRelations(DatabaseTable *table)
 {
 	QList<DatabaseRelationship *> result;
-	foreach (DatabaseRelationship *relation, m_relations) {
+	foreach (DatabaseRelationship *relation, itemsByType<DatabaseRelationship>()) {
 		if (relation->source() == table || relation->target() == table) {
 			result << relation;
 		}
@@ -208,7 +216,8 @@ DiagramDocument::save(const QString &fileName)
 	QDomElement tableListElement = doc.createElement("table-list");
 	root.appendChild(tableListElement);
 
-	foreach (DatabaseTable *table, m_tables) {
+	QList<DatabaseTable *> tables = itemsByType<DatabaseTable>();
+	foreach (DatabaseTable *table, tables) {
 		QDomElement tableElement = doc.createElement("table");
 		tableListElement.appendChild(tableElement);
 		
@@ -238,10 +247,10 @@ DiagramDocument::save(const QString &fileName)
 	QDomElement relationListElement = doc.createElement("relation-list");
 	root.appendChild(relationListElement);
 
-	foreach (DatabaseRelationship *relation, m_relations) {
+	foreach (DatabaseRelationship *relation, itemsByType<DatabaseRelationship>()) {
 		QDomElement relationElement = doc.createElement("relation");
-		relationElement.setAttribute("from", QString::number(m_tables.indexOf(relation->sourceTable())));
-		relationElement.setAttribute("to", QString::number(m_tables.indexOf(relation->targetTable())));
+		relationElement.setAttribute("from", QString::number(tables.indexOf(relation->sourceTable())));
+		relationElement.setAttribute("to", QString::number(tables.indexOf(relation->targetTable())));
 		relationListElement.appendChild(relationElement);
 	}
 
@@ -269,6 +278,8 @@ DiagramDocument::load(const QString &fileName)
 
 	QDomElement root = doc.firstChildElement("database-model");
 
+	QList<DatabaseTable *> tables;
+
 	QDomElement tableListElement = root.firstChildElement("table-list");
 	QDomElement tableElement = tableListElement.firstChildElement("table");
 	while (!tableElement.isNull()) {
@@ -297,8 +308,8 @@ DiagramDocument::load(const QString &fileName)
 			columnElement = columnElement.nextSiblingElement("column");
 		}
 
-		m_tables << table;
 		addItem(table);
+		tables.append(table);
 
 		tableElement = tableElement.nextSiblingElement("table");
 	}
@@ -313,12 +324,10 @@ DiagramDocument::load(const QString &fileName)
 			int index1 = relationElement.attribute("to").toInt(&ok);
 			if (ok) {
 				DatabaseRelationship *relation = new DatabaseRelationship();
-				relation->setSource(m_tables[index0]);
-				relation->setTarget(m_tables[index1]);
-				m_relations << relation;
+				relation->setSource(tables[index0]);
+				relation->setTarget(tables[index1]);
 				addItem(relation);
 			}
 		}
 	}
-
 }
diff --git a/src/diagramdocument.h b/src/diagramdocument.h
index c97abcb..2c87fe0 100644
--- a/src/diagramdocument.h
+++ b/src/diagramdocument.h
@@ -59,6 +59,8 @@ public:
 	void itemMoved(DiagramItem *table);
 	void itemHasMoved(DiagramItem *table);
 
+	template <class T> QList<T *> itemsByType();
+
 signals:
 	void modeChanged(DiagramDocument::Mode mode);
 
@@ -80,8 +82,6 @@ private:
 	Mode m_mode;
 	QString m_fileName;
 	QGraphicsLineItem *m_line;
-	QList<DatabaseTable *> m_tables;
-	QList<DatabaseRelationship *> m_relations;
 };
 
 #endif