Occupation system: store multiple indices for the same train, introduce reverse_lookup_sel() to select appropriate index out of multiple based on a heuristic

1 files changed, 74 insertions, 46 deletions

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