Fixing bug when un-equipping two handed weapons.

src/org/leumasjaffe/charsheet/view/inventory/EquipmentPanel.java

@@ -24,6 +24,7 @@ import java.awt.GridBagConstraints;
 import java.awt.GridBagLayout;
 import java.awt.Insets;
 import java.util.EnumSet;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 
@@ -84,39 +85,52 @@ public class EquipmentPanel extends JPanel {
 		final DDItem armor = inv.getEquipment().get(BODY);
 		if ( armor != null && armor.isArmor() ) {
 			manual.add(BODY);
-			createWithRightClickMenu(ArmorPanel::new, inv, BODY);
+			createWithRightClickMenu(ArmorPanel::new, inv, BODY, BODY);
 		}
 		
 		final DDItem main = inv.getEquipment().get(MAIN_HAND);
 		final DDItem off = inv.getEquipment().get(OFF_HAND);
-		if ( off != null && off.isArmor() ) {
-			manual.add(OFF_HAND);
-			createWithRightClickMenu(ShieldPanel::new, inv, OFF_HAND);
-		}
-		if ( main != null && main.isWeapon() ) {
-			manual.add(MAIN_HAND);
-			createWithRightClickMenu(WeaponPanel::new, inv, MAIN_HAND);
-		}
-		if ( off != null && off.isWeapon() ) {
-			manual.add(OFF_HAND);
-			if ( off != main || couldDualWieldThis(off) ) {
-				createWithRightClickMenu(WeaponPanel::new, inv, OFF_HAND);
-			}
+	
+		final Optional<Function<DDItem, JPanel>> makeMain = getEquipmentRightClickPanelFactory(main);
+		final Optional<Function<DDItem, JPanel>> makeOff = getEquipmentRightClickPanelFactory(off);
+		
+		if (main == off && !couldDualWieldThis(main)) {
+			makeMain.ifPresent(f -> {
+				manual.add(MAIN_HAND);
+				manual.add(OFF_HAND);
+				createWithRightClickMenu(f, inv, TWO_HANDS, MAIN_HAND);
+			});
+		} else {
+			makeMain.ifPresent(f -> {
+				manual.add(MAIN_HAND);
+				createWithRightClickMenu(f, inv, MAIN_HAND, MAIN_HAND);
+			});
+			makeOff.ifPresent(f -> {
+				manual.add(OFF_HAND);
+				createWithRightClickMenu(f, inv, OFF_HAND, OFF_HAND);
+			});
 		}
 		
 		inv.getEquipment().keySet().stream().filter( slot -> ! manual.contains(slot) )
 		.forEach( slot -> {
-			createWithRightClickMenu(null, inv, slot);
+			createWithRightClickMenu(null, inv, slot, slot);
 		});
 	}
 
+	private Optional<Function<DDItem, JPanel>> getEquipmentRightClickPanelFactory(final DDItem item) {
+		if (item == null) { return Optional.empty(); }
+		else if (item.isWeapon()) { return Optional.of(WeaponPanel::new); }
+		else if (item.isArmor()) { return Optional.of(ShieldPanel::new); }
+		else { return Optional.empty(); }
+	}
+
 	private boolean couldDualWieldThis(final DDItem item) {
-		return item.getSlot() == EquipmentSlot.ONE_HAND && item.getCount().value() > 1;
+		return item != null && item.getSlot() == EquipmentSlot.ONE_HAND && item.getCount().value() > 1;
 	}
 	
 	private void createWithRightClickMenu(final Function<DDItem, JPanel> make,
-			final DDInventory inv, final EquipmentSlot slot) {
-		final DDItem item = inv.getEquipment().get(slot);
+			final DDInventory inv, final EquipmentSlot slot, EquipmentSlot get) {
+		final DDItem item = inv.getEquipment().get(get);
 		final JPanel panel = make.apply(item);
 		equipment.add(panel);
 		panel.addMouseListener(new PopClickListener(