support new gpu_mem config.txt option for setting memory split

raspi-config

@@ -91,6 +91,21 @@ EOF
 mv "$3.bak" "$3"
 }
 
+get_config_var() {
+  lua - "$1" "$2" <<EOF
+local key=assert(arg[1])
+local fn=assert(arg[2])
+local file=assert(io.open(fn))
+for line in file:lines() do
+  local val = line:match("^#?%s*"..key.."=(.*)$")
+  if (val ~= nil) then
+    print(val)
+    break
+  end
+end
+EOF
+}
+
 # $1 is 0 to disable overscan, 1 to disable it
 set_overscan() {
   # Stop if /boot is not a mountpoint
@@ -141,16 +156,33 @@ do_change_timezone() {
 }
 
 do_memory_split() {
-  get_current_memory_split
-  MEMSPLIT=$(whiptail --menu "Set memory split.\n$MEMSPLIT_DESCRIPTION" 20 60 10 \
-    "240" "240MiB for ARM, 16MiB for VideoCore" \
-    "224" "224MiB for ARM, 32MiB for VideoCore" \
-    "192" "192MiB for ARM, 64MiB for VideoCore" \
-    "128" "128MiB for ARM, 128MiB for VideoCore" \
-    3>&1 1>&2 2>&3)
-  if [ $? -eq 0 ]; then
-    set_memory_split ${MEMSPLIT}
-    ASK_TO_REBOOT=1
+  if [ -e /boot/start_cd.elf ]; then
+    # New-style memory split setting
+    if ! mountpoint -q /boot; then
+      return 1
+    fi
+    ## get current memory split from /boot/config.txt
+    CUR_GPU_MEM=$(get_config_var gpu_mem /boot/config.txt)
+    [ -z "$CUR_GPU_MEM" ] && CUR_GPU_MEM=64
+    ## ask users what gpu_mem they want
+    NEW_GPU_MEM=$(whiptail --inputbox "How much memory should the GPU have?  e.g. 16/32/64/128/256" \
+      20 70 "$CUR_GPU_MEM" 3>&1 1>&2 2>&3)
+    if [ $? -eq 0 ]; then
+      set_config_var gpu_mem "$NEW_GPU_MEM" /boot/config.txt
+      ASK_TO_REBOOT=1
+    fi
+  else # Old firmware so do start.elf renaming
+    get_current_memory_split
+    MEMSPLIT=$(whiptail --menu "Set memory split.\n$MEMSPLIT_DESCRIPTION" 20 60 10 \
+      "240" "240MiB for ARM, 16MiB for VideoCore" \
+      "224" "224MiB for ARM, 32MiB for VideoCore" \
+      "192" "192MiB for ARM, 64MiB for VideoCore" \
+      "128" "128MiB for ARM, 128MiB for VideoCore" \
+      3>&1 1>&2 2>&3)
+    if [ $? -eq 0 ]; then
+      set_memory_split ${MEMSPLIT}
+      ASK_TO_REBOOT=1
+    fi
   fi
 }