4 files changed, 87 insertions, 57 deletions

diff --git a/advtrains/couple.lua b/advtrains/couple.lua
index 3e6c432..1318c12 100644
--- a/advtrains/couple.lua
+++ b/advtrains/couple.lua
@@ -79,8 +79,9 @@ function advtrains.train_check_couples(train)
 	end
 	if not train.cpl_front then
 		-- recheck front couple
-		local front_trains, pos = advtrains.occ.get_occupations(train, atround(train.index) + CPL_CHK_DST)
+		local pos = advtrains.path_get(train, atround(train.index) + CPL_CHK_DST)
 		if advtrains.is_node_loaded(pos) then -- if the position is loaded...
+			local front_trains = advtrains.occ.reverse_lookup_sel(pos, "in_train")
 			for tid, idx in pairs(front_trains) do
 				local other_train = advtrains.trains[tid]
 				if not advtrains.train_ensure_init(tid, other_train) then
@@ -109,8 +110,9 @@ function advtrains.train_check_couples(train)
 	end
 	if not train.cpl_back then
 		-- recheck back couple
-		local back_trains, pos = advtrains.occ.get_occupations(train, atround(train.end_index) - CPL_CHK_DST)
+		local pos = advtrains.path_get(train, atround(train.end_index) - CPL_CHK_DST)
 		if advtrains.is_node_loaded(pos) then -- if the position is loaded...
+			local back_trains = advtrains.occ.reverse_lookup_sel(pos, "in_train")
 			for tid, idx in pairs(back_trains) do
 				local other_train = advtrains.trains[tid]
 				if not advtrains.train_ensure_init(tid, other_train) then
diff --git a/advtrains/occupation.lua b/advtrains/occupation.lua
index db39991..6852dfa 100644
--- a/advtrains/occupation.lua
+++ b/advtrains/occupation.lua
@@ -86,9 +86,10 @@ end
 
 function o.set_item(train_id, pos, idx)
 	local t = occgetcreate(pos)
+	assert(idx)
 	local i = 1
 	while t[i] do
-		if t[i]==train_id then
+		if t[i]==train_id and t[i+1]==index then
 			break
 		end
 		i = i + 2
@@ -98,25 +99,30 @@ function o.set_item(train_id, pos, idx)
 end
 
 
-function o.clear_item(train_id, pos)
+function o.clear_all_items(train_id, pos)
 	local t = occget(pos)
 	if not t then return end
 	local i = 1
-	local moving = false
 	while t[i] do
 		if t[i]==train_id then
-			if moving then
-				-- if, for some occasion, there should be a duplicate entry, erase this one too
-				atwarn("Duplicate occupation entry at",pos,"for train",train_id,":",t)
-				i = i - 2
-			end
-			moving = true
+			table.remove(t, i)
+			table.remove(t, i)
+		else
+			i = i + 2
 		end
-		if moving then
-			t[i]   = t[i+2]
-			t[i+1] = t[i+3]
+	end
+end
+function o.clear_specific_item(train_id, pos, index)
+	local t = occget(pos)
+	if not t then return end
+	local i = 1
+	while t[i] do
+		if t[i]==train_id and t[i+1]==index then
+			table.remove(t, i)
+			table.remove(t, i)
+		else
+			i = i + 2
 		end
-		i = i + 2
 	end
 end
 
@@ -143,64 +149,86 @@ function o.check_collision(pos, train_id)
 	return false
 end
 
--- Gets a mapping of train id's to indexes of trains that share this path item with this train
--- The train itself will not be included.
--- If the requested index position is off-track, returns {}.
--- returns (table with train_id->index), position
-function o.get_occupations(train, index)
-	local ppos, ontrack = advtrains.path_get(train, index)
-	if not ontrack then
-		atlog("Train",train.id,"get_occupations requested off-track",index)
-		return {}, ppos
-	end
+-- Gets a mapping of train id's to indexes of trains that have a path item at this position
+-- Note that the case where 2 or more indices are at a position only occurs if there is a track loop.
+-- returns (table with train_id->{index1, index2...})
+function o.reverse_lookup(ppos)
 	local pos = advtrains.round_vector_floor_y(ppos)
 	local t = occget(pos)
 	if not t then return {} end
 	local r = {}
 	local i = 1
-	local train_id = train.id
 	while t[i] do
 		if t[i]~=train_id then
-			r[t[i]] = t[i+1]
+			if not r[t[i]] then r[t[i]] = {} end
+			table.insert(r[t[i]], t[i+1])
 		end
 		i = i + 2
 	end
-	return r, pos
+	return r
 end
--- Gets a mapping of train id's to indexes of trains that stand or drive over
+
+-- Gets a mapping of train id's to indexes of trains that have a path item at this position.
+-- Quick variant: will only return one index per train (the latest one added)
 -- returns (table with train_id->index)
-function o.get_trains_at(ppos)
+function o.reverse_lookup_quick(ppos)
 	local pos = advtrains.round_vector_floor_y(ppos)
 	local t = occget(pos)
 	if not t then return {} end
 	local r = {}
 	local i = 1
 	while t[i] do
-		local train = advtrains.trains[t[i]]
-		local idx = t[i+1]
-		if train.end_index - 0.5 <= idx and idx <= train.index + 0.5 then
-			r[t[i]] = idx
-		end
+		r[t[i]] = t[i+1]
 		i = i + 2
 	end
 	return r
 end
 
--- Gets a mapping of train id's to indexes of trains that have a path
--- generated over this node
--- returns (table with train_id->index)
-function o.get_trains_over(ppos)
-	local pos = advtrains.round_vector_floor_y(ppos)
-	local t = occget(pos)
-	if not t then return {} end
+local OCC_CLOSE_PROXIMITY = 3
+-- Gets a mapping of train id's to index of trains that have a path item at this position. Selects at most one index based on a given heuristic, or even none if it does not match the heuristic criterion
+-- returns (table with train_id->index), position
+-- "in_train": first index that lies between train index and end index
+-- "first_ahead": smallest index that is > current index
+-- "before_end"(default): smallest index that is > end index
+-- "close_proximity": within 3 indices close to the train index and end_index
+-- "any": just output the first index found and do not check further (also occurs if both "in_train" and "first_ahead" heuristics have failed
+function o.reverse_lookup_sel(pos, heuristic)
+	if not heuristic then heuristic = "before_end" end
+	local om = o.reverse_lookup(pos)
 	local r = {}
-	local i = 1
-	while t[i] do
-		local idx = t[i+1]
-		r[t[i]] = idx
-		i = i + 2
+	for tid, idxs in pairs(om) do
+		r[tid] = idxs[1]
+		if heuristic~="any" then
+			--must run a heuristic
+			--atdebug("reverse_lookup_sel is running heuristic for", pos,heuristic,"idxs",table.concat(idxs,","))
+			local otrn = advtrains.trains[tid]
+			advtrains.train_ensure_init(tid, otrn)
+			local h_value
+			for _,idx in ipairs(idxs) do
+				if heuristic == "first_ahead" and idx > otrn.index and (not h_value or h_value>idx) then
+					h_value = idx
+				end
+				if heuristic == "before_end" and idx > otrn.end_index and (not h_value or h_value>idx) then
+					h_value = idx
+				end
+				if heuristic == "in_train" and idx < otrn.index and idx > otrn.end_index then
+					h_value = idx
+				end
+				if heuristic == "close_proximity" and idx < (otrn.index + OCC_CLOSE_PROXIMITY) and idx > (otrn.end_index - OCC_CLOSE_PROXIMITY) then
+					h_value = idx
+				end
+			end
+			r[tid] = h_value
+			--atdebug(h_value,"chosen")
+		end
 	end
-	return r
+	return r, pos
+end
+-- Gets a mapping of train id's to indexes of trains that stand or drive over
+-- returns (table with train_id->index)
+function o.get_trains_at(ppos)
+	local pos = advtrains.round_vector_floor_y(ppos)
+	return o.reverse_lookup_sel(pos, "in_train")
 end
 
 advtrains.occ = o
diff --git a/advtrains/path.lua b/advtrains/path.lua
index 19387b1..15a61fe 100644
--- a/advtrains/path.lua
+++ b/advtrains/path.lua
@@ -119,7 +119,7 @@ function advtrains.path_invalidate(train, ignore_lock)
 
 	if train.path then
 		for i,p in pairs(train.path) do
-			advtrains.occ.clear_item(train.id, advtrains.round_vector_floor_y(p))
+			advtrains.occ.clear_all_items(train.id, advtrains.round_vector_floor_y(p))
 		end
 	end
 	train.path = nil
@@ -393,7 +393,7 @@ local PATH_CLEAR_KEEP = 4
 function advtrains.path_clear_unused(train)
 	local i
 	for i = train.path_ext_b, train.path_req_b - PATH_CLEAR_KEEP do
-		advtrains.occ.clear_item(train.id, advtrains.round_vector_floor_y(train.path[i]))
+		advtrains.occ.clear_specific_item(train.id, advtrains.round_vector_floor_y(train.path[i]), i)
 		train.path[i] = nil
 		train.path_dist[i-1] = nil
 		train.path_cp[i] = nil
@@ -434,18 +434,19 @@ end
 -- Projects the path of "train" onto the path of "onto_train_id", and returns the index on onto_train's path
 -- that corresponds to "index" on "train"'s path, as well as whether both trains face each other
 -- index may be fractional
+-- heuristic: see advtrains.occ.reverse_lookup_sel()
 -- returns: res_index, trains_facing
 -- returns nil when path can not be projected, either because trains are on different tracks or
 -- node at "index" happens to be on a turnout and it's the wrong direction
 -- Note - duplicate with similar functionality is in train_step_b() - that code combines train detection with projecting
-function advtrains.path_project(train, index, onto_train_id)
+function advtrains.path_project(train, index, onto_train_id, heuristic)
 	local base_idx = atfloor(index)
 	local frac_part = index - base_idx
 	local base_pos = advtrains.path_get(train, base_idx)
 	local base_cn =  train.path_cn[base_idx]
 	local otrn = advtrains.trains[onto_train_id]
 	-- query occupation
-	local occ = advtrains.occ.get_trains_over(base_pos)
+	local occ = advtrains.occ.reverse_lookup_sel(base_pos, heuristic)
 	-- is wanted train id contained?
 	local ob_idx = occ[onto_train_id]
 	if not ob_idx then
diff --git a/advtrains/trainlogic.lua b/advtrains/trainlogic.lua
index 288e224..f136577 100644
--- a/advtrains/trainlogic.lua
+++ b/advtrains/trainlogic.lua
@@ -619,7 +619,7 @@ function advtrains.train_step_b(id, train, dtime)
 		local base_cn =  train.path_cn[base_idx]
 		--atdebug(id,"Begin Checking for on-track collisions new_idx=",new_index_curr_tv,"base_idx=",base_idx,"base_pos=",base_pos,"base_cn=",base_cn)
 		-- query occupation
-		local occ = advtrains.occ.get_trains_over(base_pos)
+		local occ = advtrains.occ.reverse_lookup_sel(base_pos, "close_proximity")
 		-- iterate other trains
 		for otid, ob_idx in pairs(occ) do
 			if otid ~= id then
@@ -649,7 +649,7 @@ function advtrains.train_step_b(id, train, dtime)
 
 				-- Phase 2 - project ref_index back onto our path and check again (necessary because there might be a turnout on the way and we are driving into the flank
 				if target_is_inside then
-					local our_index = advtrains.path_project(otrn, ref_index, id)
+					local our_index = advtrains.path_project(otrn, ref_index, id, "before_end")
 					--atdebug("Backprojected our_index",our_index)
 					if our_index and our_index <= new_index_curr_tv
 							and our_index >= train.index then --FIX: If train was already past the collision point in the previous step, there is no collision! Fixes bug with split_at_index
@@ -1218,7 +1218,6 @@ function advtrains.invert_train(train_id)
 	advtrains.update_trainpart_properties(train_id, true)
 	
 	-- recalculate path
-	advtrains.train_ensure_init(train_id, train)
 	
 	-- If interlocking present, check whether this train is in a section and then set as shunt move after reversion
 	if advtrains.interlocking and train.il_sections and #train.il_sections > 0 then
@@ -1244,7 +1243,7 @@ function advtrains.invalidate_all_paths(pos)
 	local tab
 	if pos then
 		-- if position given, check occupation system
-		tab = advtrains.occ.get_trains_over(pos)
+		tab = advtrains.occ.reverse_lookup_quick(pos)
 	else
 		tab = advtrains.trains
 	end
@@ -1257,7 +1256,7 @@ end
 -- Calls invalidate_path_ahead on all trains occupying (having paths over) this node
 -- Can be called during train step.
 function advtrains.invalidate_all_paths_ahead(pos)
-	local tab = advtrains.occ.get_trains_over(pos)
+	local tab = advtrains.occ.reverse_lookup_sel(pos, "first_ahead")
 	
 	for id,index in pairs(tab) do
 		local train = advtrains.trains[id]