From 2365f7ec8da6b7bb69323dab6f0a3dc462f90761 Mon Sep 17 00:00:00 2001
From: delucis <swithinbank@gmail.com>
Date: Thu, 9 Jun 2016 01:22:14 +0200
Subject: [PATCH] Rewrite assign_menu_location().

---
 functions.php | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/functions.php b/functions.php
index c2e38e7..9d497ec 100644
--- a/functions.php
+++ b/functions.php
@@ -25,16 +25,10 @@ function register_my_menu() {
 
 function assign_menu_location() {
   $locations = get_nav_menu_locations();
-  if(!empty($locations)) {
-    foreach($locations as $locationId => $menuValue) {
-      switch($locationId) {
-        case 'primary':
-            $menu = get_term_by('slug', 'main-menu', 'nav_menu');
-        break;
-      }
-      if(isset($menu)) {
-        $locations[$locationId] = $menu->term_id;
-      }
+  if(empty($locations) || $locations['primary'] == 0) {
+    $menu = get_term_by('slug', 'main-menu', 'nav_menu');
+    if(isset($menu)) {
+      $locations['primary'] = $menu->term_id;
     }
     set_theme_mod('nav_menu_locations', $locations);
   }