Replace the code that guarantees the device stays in direct mode
with iio_device_claim_direct_mode() which does same.  Includes a
tiny bit of refactoring (single case -> if) and simplified return
flow.

Also prevents a theoretical race condition by effectively taking the lock
a little earlier than before.
Signed-off-by: Alison Schofield <amsfield22@gmail.com>
Cc: Daniel Baluta <daniel.baluta@gmail.com>
Acked-by: Jacek Anaszewski <j.anaszewski@samsung.com>
Signed-off-by: Jonathan Cameron <jic23@kernel.org>