Fix check.sh to deal with single-line 'vars' definition

The sed script replacing vars with proc_vars incorrectly handled
single-line vars definition.

Signed-off-by: Catalin Marinas <catalin.marinas@arm.com>
diff --git a/check.sh b/check.sh
index 06e77e4..b7a8033 100755
--- a/check.sh
+++ b/check.sh
@@ -20,6 +20,9 @@
 fi
 
 # Split << pc, stack >> out of the default vars and generate proc_vars
-sed -i -e "/^vars\s*==/,/>>/{s/\<vars\>/proc_vars/;s/\<pc,\s*\|\<stack,\s*//g;s/>>/>>\n\nvars == << proc_vars, pc, stack >>/}" $SPEC.tla
+# Match single line and multiline patterns
+SUBST="{s/\<vars\>/proc_vars/;s/\<pc,\s*\|\<stack,\s*//g;s/>>/>>\n     vars == << proc_vars, pc, stack >>/}"
+sed -i -e "/^vars\s*==.*>>$/$SUBST" $SPEC.tla
+sed -i -e "/^vars\s*==[^>]*/,/>>$/$SUBST" $SPEC.tla
 
 tlc -workers $(nproc) $@ $SPEC.tla | tee -a $SPEC.log