Fix bugs in buttons settings menu

src/buttons.js

@@ -1,6 +1,16 @@
 function Buttons(session) {
 	var mapping;
 	var mappingsStack = [];
+	var buttons = {
+		'A': 0,
+		'B': 1,
+		'Left': 2,
+		'Right': 3,
+		'Up': 4,
+		'Down': 5,
+		'Start': 6,
+		'Select': 7
+	};
 
 	document.addEventListener('keydown', function(event) {
 		var button = session.getOptions().layout[event.key];
@@ -11,19 +21,28 @@ function Buttons(session) {
 	});
 
 	return {
-		findKey: findKey,
+		assign: assign,
+		code: code,
 		map: map,
 		pop: pop,
 		push: push
 	};
 
-	function findKey(button) {
+	function assign(key, button) {
+		var layoutDiff = {};
 		var layout = session.getOptions().layout;
-		for(var key in layout) {
-			if(layout[key] == button) {
-				return key;
+		for(var currentKey in layout) {
+			if(layout[currentKey] == button) {
+				layoutDiff[currentKey] = undefined;
+				break;
 			}
 		}
+		layoutDiff[key] = button;
+		return {layoutDiff: layoutDiff, button: layout[key] != button ? layout[key] : null};
+	}
+
+	function code(button) {
+		return buttons[button];
 	}
 
 	function map(newMapping) {

src/ui.js

@@ -204,22 +204,22 @@ function UI(async, buttons, dom, screen, session) {
 
 	function setLayout() {
 		var layout = session.getOptions().layout
-		var buttons = {
+		var buttonsMenu = {
 			cancel: function() {},
 			entries: [],
 			name: "layout"
 		};
 		for(var key in layout) {
-			buttons.entries.push({
+			buttonsMenu.entries[buttons.code(layout[key])] = {
 				label: keyPromptLabel(key, layout[key]),
 				action: setKey(layout[key])
-			});
+			};
 		}
-		menu(buttons);
+		menu(buttonsMenu);
 	}
 
 	function setKey(button) {
-		return function() {
+		return function(f) {
 			buttons.push({});
 			var elem = dom.make('div', {
 				class: ['framed', 'center'],
@@ -227,14 +227,19 @@ function UI(async, buttons, dom, screen, session) {
 			});
 			screen.currentMenu().getElementsByClassName('selected')[0].textContent = keyPromptLabel(null, button);
 			var promptKey = function(event) {
-				var layout = {};
-				layout[event.key] = button;
-				layout[buttons.findKey(button)] = undefined;
-				session.setOptions({layout: layout});
+				var assignment = buttons.assign(event.key, button);
+				session.setOptions({layout: assignment.layoutDiff});
 				screen.clear(elem);
 				screen.currentMenu().getElementsByClassName('selected')[0].textContent = keyPromptLabel(event.key, button);
 				document.removeEventListener('keydown', promptKey);
 				buttons.pop();
+				if(assignment.button != undefined) {
+					screen.select(buttons.code(assignment.button));
+					return setKey(assignment.button)(function () {
+						screen.select(buttons.code(button));
+					});
+				}
+				return f();
 			};
 			screen.show(elem);
 			document.addEventListener('keydown', promptKey);