better type for Retriever

Putting a callback in the Retriever type allows for the callback to
remove the retrieved file when it's done with it.

I did not really want to make Retriever be fixed to Annex Bool,
but when I tried to use Annex a, I got into some type of type mess.

2 files changed, 48 insertions, 14 deletions

diff --git a/Remote/Helper/Chunked.hs b/Remote/Helper/Chunked.hs
index 102ced8f4..ae949abc3 100644
--- a/Remote/Helper/Chunked.hs
+++ b/Remote/Helper/Chunked.hs
@@ -249,26 +249,28 @@ retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
 		let p = maybe basep
 			(offsetMeterUpdate basep . toBytesProcessed)
 			offset
-		v <- tryNonAsyncAnnex $ retriever (encryptor k) p
-		case v of
-			Left e
-				| null ls -> giveup e
-				| otherwise -> firstavail currsize ls
-			Right content -> do
+		v <- tryNonAsyncAnnex $
+			retriever (encryptor k) p $ \content ->
 				bracketIO (maybe opennew openresume offset) hClose $ \h -> do
 					tosink h p content
 					let sz = toBytesProcessed $
 						fromMaybe 0 $ keyChunkSize k
 					getrest p h sz sz ks
+						`catchNonAsyncAnnex` giveup
+		case v of
+			Left e
+				| null ls -> giveup e
+				| otherwise -> firstavail currsize ls
+			Right r -> return r
 
 	getrest _ _ _ _ [] = return True
 	getrest p h sz bytesprocessed (k:ks) = do
 		let p' = offsetMeterUpdate p bytesprocessed
-		tosink h p' =<< retriever (encryptor k) p'
+		retriever (encryptor k) p' $ tosink h p'
 		getrest p h sz (addBytesProcessed bytesprocessed sz) ks
 
 	getunchunked = bracketIO opennew hClose $ \h -> do
-		tosink h basep =<< retriever (encryptor basek) basep
+		retriever (encryptor basek) basep $ tosink h basep
 		return True
 
 	opennew = openBinaryFile dest WriteMode
@@ -288,10 +290,13 @@ retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
 	 - it is not responsible for updating progress (often it cannot).
 	 - Instead, the sink is passed a meter to update  as it consumes
 	 - the ByteString. -}
-	tosink h p (ByteContent b) = liftIO $
+	tosink h p (ByteContent b) = liftIO $ do
 		sink h (Just p) b
-	tosink h _ (FileContent f) = liftIO $
+		return True
+	tosink h _ (FileContent f) = liftIO $ do
 		sink h Nothing =<< L.readFile f
+		nukeFile h
+		return True
 
 {- Can resume when the chunk's offset is at or before the end of
  - the dest file. -}
diff --git a/Remote/Helper/ChunkedEncryptable.hs b/Remote/Helper/ChunkedEncryptable.hs
index 024a53309..550a6934b 100644
--- a/Remote/Helper/ChunkedEncryptable.hs
+++ b/Remote/Helper/ChunkedEncryptable.hs
@@ -6,8 +6,6 @@
  -}
 
 {-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE ImpredicativeTypes #-}
-{-# LANGUAGE FlexibleContexts #-}
 
 module Remote.Helper.ChunkedEncryptable (
 	Preparer,
@@ -39,17 +37,48 @@ import Annex.Exception
 
 import qualified Data.ByteString.Lazy as L
 
+-- Use when nothing needs to be done to prepare a helper.
 simplyPrepare :: helper -> Preparer helper
 simplyPrepare helper _ a = a $ Just helper
 
+-- Use to run a check when preparing a helper.
 checkPrepare :: (Key -> Annex Bool) -> helper -> Preparer helper
 checkPrepare checker helper k a = ifM (checker k)
 	( a (Just helper)
 	, a Nothing
 	)
 
-{- Modifies a base Remote to support both chunking and encryption.
- -}
+-- A Storer that expects to be provided with a file containing
+-- the content of the key to store.
+fileStorer :: (Key -> FilePath -> MeterUpdate -> Annex Bool) -> Storer
+fileStorer a k (FileContent f) m = a k f m
+fileStorer a k (ByteContent b) m = withTmp k $ \f -> do
+	liftIO $ L.writeFile f b
+	a k f m
+
+-- A Storer that expects to be provided with a L.ByteString of
+-- the content to store.
+byteStorer :: (Key -> L.ByteString -> MeterUpdate -> Annex Bool) -> Storer
+byteStorer a k c m = withBytes c $ \b -> a k b m
+
+-- A Retriever that writes the content of a Key to a provided file.
+-- It is responsible for updating the progress meter as it retrieves data.
+fileRetriever :: (FilePath -> Key -> MeterUpdate -> Annex ()) -> Retriever
+fileRetriever a k m callback = bracketAnnex (prepTmp k) (liftIO . nukeFile) go
+  where
+	go f = do
+		a f k m
+		callback (FileContent f)
+
+-- A Retriever that generates a L.ByteString containing the Key's content.
+byteRetriever :: (Key -> Annex L.ByteString) -> Retriever
+byteRetriever a k _m callback = callback =<< (ByteContent <$> a k)
+
+withBytes :: ContentSource -> (L.ByteString -> Annex a) -> Annex a
+withBytes (ByteContent b) a = a b
+withBytes (FileContent f) a = a =<< liftIO (L.readFile f)
+
+-- Modifies a base Remote to support both chunking and encryption.
 chunkedEncryptableRemote
 	:: RemoteConfig
 	-> Preparer Storer