#!/usr/bin/env sh
# (C) Martin V\"ath <martin@mvath.de>"

Echo() {
	printf '%s\n' "${*}"
}

Echon() {
	printf '%s' "${*}"
}

Warning() {
	Echo "${0##*/}: ${*}" >&2
	exit 1
}

Fatal() {
	Warning "${*}"
	exit 1
}

Usage() {
	Echo "Usage: ${0##*/} [options] SIZE [DIR]
Prepare a zram device and use it as swap (resp. mount it under DIR).
SIZE is the maximal size in megabytes.
For umounting/freeing the zram device, use SIZE=0.
The following options are available.
An empty argument means the same as if the option is not specified.
-d DEV   Use zram device DEV. If not specified, DEV=0 is assumed.
         Make sure to use the matching value for umounting (SIZE=0)!
-c OWNER If specified, chown to OWNER (resp. OWNER:GROUP) after mounting.
         If not specified, the default owner (root:root) is not changed.
-m MODE  If specified, chmod DIR to MODE after mounting.
-o OPTS  If specified, mount DIR with -o OPTS.
-p PRIO  Use priority PRIO for the swap device.
         If not specified, PRIO=16383 is assumed.
         Use PRIO=- if you want to keep the default priority (-1).
-t TYPE  Use a filesystem of type TYPE if DIR is specified.
         TYPE can be either ext2 or ext4.
         If not specified, TYPE=ext4 is assumed.
-k       Do no attempt to umount/free a previously used zram under this device."
	exit ${1:-1}
}

dev=
prio=
fstype=
opts=
mode=
owgr=
keep=false
while getopts 'c:d:m:o:p:t:khH?' opt
do	case ${opt} in
	c)	owgr=${OPTARG};;
	d)	dev=${OPTARG};;
	m)	mode=${OPTARG};;
	o)	opts=${OPTARG};;
	p)	prio=${OPTARG};;
	t)	fstype=${OPTARG};;
	k)	keep=:;;
	*)	Usage 0;;
	esac
done
shift $(( ${OPTIND} -1 ))
[ ${#} -eq 1 ] || [ ${#} -eq 2 ] || Usage

# Defaults for empty/unspecified options:
: ${dev:=0} ${prio:=16383} ${fstype:=ext4}

umount=:
[ "${1:-0}" -ge 0 ] || Usage
if [ "${1:-0}" -ne 0 ]
then	umount=false
	size=$(( ${1} * 1024 * 1024 )) && [ "${size}" -gt 0 ] || Usage
fi

swap=:
if [ ${#} -eq 2 ]
then	swap=false
	dir=${2%/}
	if [ -z "${dir}" ] || ! test -d "${dir}"
	then	if ${umount}
		then	Warning "${dir:-(empty)} is not a directory. Umounting anyway"
		else	Fatal "${dir:-(empty)} is not a directory"
		fi
	fi
fi

devnode="/dev/zram${dev}"
if ! test -e "${devnode}"
then	${umount} && exit
	modprobe zram >/dev/null 2>&1
	test -e "${devnode}" || Fatal "cannot find ${devnode}"
fi
block="/sys/block/zram${dev}"
test -d "${block}" || Fatal "cannot find ${block}"

if ! ${keep}
then	swapoff -- "${devnode}" >/dev/null 2>&1 \
	|| umount -- "${devnode}" >/dev/null 2>&1 \
	&& Echon 1 >"${block}/reset"
fi
status=${?}
${umount} && exit ${status}

Echon "${size}" >"${block}/disksize" || Fatal "cannot set zram${dev} size"

if ${swap}
then	mkswap "${devnode}" >/dev/null || Fatal "mkswap ${devnode} failed"
	if [ "${prio}" = '-' ]
	then	swapon -- "${devnode}"
	else	swapon -p "${prio}" -- "${devnode}"
	fi || Fatal "swapon failed"
	exit 0
fi

fsopts='-O^huge_file,sparse_super'
case ${fstype} in
ext2)
	mkfs.ext2 -m0 "${fsopts}" "${devnode}" >/dev/null \
		|| Fatal "mkfs.ext2 ${devnode} failed"
	tune2fs -c0 -i0 -m0 "${devnode}" >/dev/null
	mount -t ext2 ${opts:+-o "${opts}"} -- "${devnode}" "${dir}" \
		|| Fatal "mount ${devnode} failed";;
ext4)
	fsopts=${fsopts}',extent,^uninit_bg,dir_nlink,extra_isize,^has_journal'
	mkfs.ext4 -m0 "${fsopts}" "${devnode}" >/dev/null \
		|| Fatal "mkfs.ext4 ${devnode} failed"
	tune2fs -c0 -i0 -m0 "${devnode}" >/dev/null
	mount -t ext4 ${opts:+-o "${opts}"} -- "${devnode}" "${dir}" \
		|| Fatal "mount ${devnode} failed";;
*)
	Fatal "unsupported filesystem ${fstype}";;
esac

# Change owner/mode if requested
[ -z "${owgr}" ] || chown -- "${owgr}" "${dir}"
[ -z "${mode}" ] || chmod -- "${mode}" "${dir}"
