6 files changed, 41 insertions, 13 deletions

diff --git a/divelist.c b/divelist.c
index d57188a68..773f52713 100644
--- a/divelist.c
+++ b/divelist.c
@@ -341,6 +341,19 @@ int get_divenr(struct dive *dive)
 	return -1;
 }
 
+int get_divesite_idx(struct dive_site *ds)
+{
+	int i;
+	struct dive_site *d;
+	// tempting as it may be, don't die when called with dive=NULL
+	if (ds)
+		for_each_dive_site(i, d) {
+			if (d->uuid == ds->uuid) // don't compare pointers, we could be passing in a copy of the dive
+				return i;
+		}
+	return -1;
+}
+
 static struct gasmix air = { .o2.permille = O2_IN_AIR, .he.permille = 0 };
 
 /* take into account previous dives until there is a 48h gap between dives */
diff --git a/divelist.h b/divelist.h
index 79e1f6cbc..91318c3b1 100644
--- a/divelist.h
+++ b/divelist.h
@@ -22,6 +22,7 @@ extern dive_trip_t *find_trip_by_idx(int idx);
 extern int trip_has_selected_dives(dive_trip_t *trip);
 extern void get_dive_gas(struct dive *dive, int *o2_p, int *he_p, int *o2low_p);
 extern int get_divenr(struct dive *dive);
+extern int get_divesite_idx(struct dive_site *ds);
 extern dive_trip_t *find_matching_trip(timestamp_t when);
 extern void remove_dive_from_trip(struct dive *dive, short was_autogen);
 extern dive_trip_t *create_and_hookup_trip_from_dive(struct dive *dive);
diff --git a/qt-models/divelocationmodel.cpp b/qt-models/divelocationmodel.cpp
index 42f283f1e..edf15b707 100644
--- a/qt-models/divelocationmodel.cpp
+++ b/qt-models/divelocationmodel.cpp
@@ -39,16 +39,10 @@ QVariant LocationInformationModel::data(const QModelIndex &index, int role) cons
 
 void LocationInformationModel::update()
 {
-	if (rowCount()) {
-		beginRemoveRows(QModelIndex(), 0, rowCount()-1);
-		endRemoveRows();
-	}
-	if (dive_site_table.nr) {
-		beginInsertRows(QModelIndex(), 0, dive_site_table.nr);
-		internalRowCount = dive_site_table.nr;
-		std::sort(dive_site_table.dive_sites, dive_site_table.dive_sites + dive_site_table.nr, dive_site_less_than);
-		endInsertRows();
-	}
+	beginResetModel();
+	internalRowCount = dive_site_table.nr;
+	std::sort(dive_site_table.dive_sites, dive_site_table.dive_sites + dive_site_table.nr, dive_site_less_than);
+	endResetModel();
 }
 
 int32_t LocationInformationModel::addDiveSite(const QString& name, int lon, int lat)
@@ -76,3 +70,14 @@ bool LocationInformationModel::setData(const QModelIndex &index, const QVariant
 	emit dataChanged(index, index);
 	return true;
 }
+
+bool LocationInformationModel::removeRows(int row, int count, const QModelIndex & parent) {
+	if(row >= rowCount())
+		return false;
+
+	beginRemoveRows(QModelIndex(), row, row);
+	struct dive_site *ds = get_dive_site(row);
+	delete_dive_site(ds->uuid);
+	endRemoveRows();
+	return true;
+}
diff --git a/qt-models/divelocationmodel.h b/qt-models/divelocationmodel.h
index 7ca0c926a..4e5adf104 100644
--- a/qt-models/divelocationmodel.h
+++ b/qt-models/divelocationmodel.h
@@ -13,6 +13,8 @@ public:
 	QVariant data(const QModelIndex &index = QModelIndex(), int role = Qt::DisplayRole) const;
 	int32_t addDiveSite(const QString& name, int lat = 0, int lon = 0);
 	bool setData(const QModelIndex &index, const QVariant &value, int role);
+	bool removeRows(int row, int count, const QModelIndex & parent = QModelIndex());
+
 public slots:
 	void update();
 private:
diff --git a/qt-ui/locationinformation.cpp b/qt-ui/locationinformation.cpp
index 75faf87d6..cc6405947 100644
--- a/qt-ui/locationinformation.cpp
+++ b/qt-ui/locationinformation.cpp
@@ -91,7 +91,7 @@ void LocationInformationWidget::acceptChanges()
 		currentDs->notes = copy_string(uiString);
 	}
 	if (dive_site_is_empty(currentDs)) {
-		delete_dive_site(currentDs->uuid);
+		LocationInformationModel::instance()->removeRow(get_divesite_idx(currentDs));
 		displayed_dive.dive_site_uuid = 0;
 	}
 
@@ -104,7 +104,7 @@ void LocationInformationWidget::acceptChanges()
 void LocationInformationWidget::rejectChanges()
 {
 	if (currentDs && dive_site_is_empty(currentDs)) {
-		delete_dive_site(currentDs->uuid);
+		LocationInformationModel::instance()->removeRow(get_divesite_idx(currentDs));
 		displayed_dive.dive_site_uuid = 0;
 	}
 	resetState();
diff --git a/qt-ui/maintab.cpp b/qt-ui/maintab.cpp
index cc09ed723..85d45b357 100644
--- a/qt-ui/maintab.cpp
+++ b/qt-ui/maintab.cpp
@@ -243,6 +243,10 @@ void MainTab::disableGeoLookupEdition()
 }
 
 void MainTab::prepareDiveSiteEdit() {
+	// TODO: This is wrong. We can only set this if we Accepted the dive site edit
+	// And not if we cancelled. Currently we are seting directly without even
+	// thinking - but too tired, fix this tomorrow.
+
 	uint32_t dive_site_uuid = LocationInformationModel::instance()->addDiveSite(tr("Unnamed"));
 	displayed_dive.dive_site_uuid = dive_site_uuid;
 	emit requestDiveSiteEdit(dive_site_uuid);
@@ -429,7 +433,10 @@ bool MainTab::isEditing()
 
 void MainTab::showLocation()
 {
-	ui.location->setCurrentText(get_dive_location(&displayed_dive));
+	if (get_dive_site_by_uuid(displayed_dive.dive_site_uuid))
+		ui.location->setCurrentText(get_dive_location(&displayed_dive));
+	else
+		ui.location->setCurrentIndex(-1);
 }
 
 void MainTab::updateDiveInfo(bool clear)