Fix frame switching after frame redirect/creation

2 files changed, 37 insertions, 9 deletions

diff --git a/background_scripts/main.coffee b/background_scripts/main.coffee
index b40907fb..4dd12420 100644
--- a/background_scripts/main.coffee
+++ b/background_scripts/main.coffee
@@ -285,11 +285,11 @@ BackgroundCommands =
   nextFrame: (count) ->
     chrome.tabs.getSelected(null, (tab) ->
       frames = framesForTab[tab.id].frames
-      currIndex = getCurrFrameIndex(frames)
+      currIndex = frames.frameIndex or getCurrFrameIndex(frames)
 
       # TODO: Skip the "top" frame (which doesn't actually have a <frame> tag),
       # since it exists only to contain the other frames.
-      newIndex = (currIndex + count) % frames.length
+      frames.frameIndex = newIndex = (currIndex + count) % frames.length
 
       chrome.tabs.sendMessage(tab.id, { name: "focusFrame", frameId: frames[newIndex].id, highlight: true }))
 
@@ -394,7 +394,7 @@ chrome.tabs.onUpdated.addListener (tabId, changeInfo, tab) ->
     code: Settings.get("userDefinedLinkHintCss")
     runAt: "document_start"
   chrome.tabs.insertCSS tabId, cssConf, -> chrome.runtime.lastError
-  updateOpenTabs(tab)
+  updateOpenTabs(tab) if "url" in changeInfo
   updateActiveState(tabId)
 
 chrome.tabs.onAttached.addListener (tabId, attachedInfo) ->
@@ -608,11 +608,30 @@ registerFrame = (request, sender) ->
 
   if (request.is_top)
     focusedFrame = request.frameId
-    framesForTab[sender.tab.id].total = request.total
+    framesForTab[sender.tab.id].frames.frameIndex = getCurrFrameIndex(framesForTab[sender.tab.id])
 
   framesForTab[sender.tab.id].frames.push({ id: request.frameId })
 
-handleFrameFocused = (request, sender) -> focusedFrame = request.frameId
+unregisterFrame = (request, sender) ->
+  return unless framesForTab[sender.tab.id]
+
+  if request.is_top # The whole tab is closing, so we can drop the frames list.
+    updateOpenTabs(sender.tab)
+    return
+
+  frames = framesForTab[sender.tab.id].frames
+
+  for index, tabDetails of frames
+    if (tabDetails.id == request.frameId)
+      frames.splice(index, 1)
+
+      frames.frameIndex-- if frames.frameIndex >= index # Adjust index if it is shifted.
+      nextFrame(0) if frames.frameIndex == index # Focus another frame if the closed one was focused.
+      return
+
+handleFrameFocused = (request, sender) ->
+  focusedFrame = request.frameId
+  framesForTab[sender.tab.id].frames.frameIndex = getCurrFrameIndex(framesForTab[sender.tab.id])
 
 getCurrFrameIndex = (frames) ->
   for i in [0...frames.length]
@@ -633,6 +652,7 @@ sendRequestHandlers =
   openUrlInCurrentTab: openUrlInCurrentTab,
   openOptionsPageInNewTab: openOptionsPageInNewTab,
   registerFrame: registerFrame,
+  unregisterFrame: unregisterFrame,
   frameFocused: handleFrameFocused,
   upgradeNotificationClosed: upgradeNotificationClosed,
   updateScrollPosition: handleUpdateScrollPosition,
@@ -640,7 +660,7 @@ sendRequestHandlers =
   isEnabledForUrl: isEnabledForUrl,
   saveHelpDialogSettings: saveHelpDialogSettings,
   selectSpecificTab: selectSpecificTab,
-  refreshCompleter: refreshCompleter
+  refreshCompleter: refreshCompleter,
   createMark: Marks.create.bind(Marks),
   gotoMark: Marks.goto.bind(Marks)
 
diff --git a/content_scripts/vimium_frontend.coffee b/content_scripts/vimium_frontend.coffee
index 63487440..8b32c98b 100644
--- a/content_scripts/vimium_frontend.coffee
+++ b/content_scripts/vimium_frontend.coffee
@@ -179,20 +179,27 @@ window.addEventListener "focus", ->
 # Initialization tasks that must wait for the document to be ready.
 #
 initializeOnDomReady = ->
-  registerFrame(window.top == window.self)
+  registerFrame()
 
   enterInsertModeIfElementIsFocused() if isEnabledForUrl
 
   # Tell the background page we're in the dom ready state.
   chrome.runtime.connect({ name: "domReady" })
 
-registerFrame = (is_top) ->
+registerFrame = ->
   chrome.runtime.sendMessage(
     handler: "registerFrame"
     frameId: frameId
-    is_top: is_top
+    is_top: window.top == window.self
     total: frames.length + 1)
 
+# Unregister the frame if we're going to exit.
+unregisterFrame = ->
+  chrome.runtime.sendMessage(
+    handler: "unregisterFrame"
+    frameId: frameId
+    is_top: window.top == window.self)
+
 #
 # Enters insert mode if the currently focused element in the DOM is focusable.
 #
@@ -1056,6 +1063,7 @@ Tween =
 
 initializePreDomReady()
 window.addEventListener("DOMContentLoaded", initializeOnDomReady)
+window.addEventListener("unload", unregisterFrame)
 
 window.onbeforeunload = ->
   chrome.runtime.sendMessage(