-- Copyright 2020 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

{-# LANGUAGE TypeApplications #-}
-- | Bind to the rubberband library.
module Util.Audio.Rubberband (
    Config(..)
    , config
    , Option -- ProcessOffline and ProcessRealTime omitted
        ( StretchElastic, StretchPrecise, TransientsCrisp
        , TransientsMixed, TransientsSmooth
        , DetectorCompound, DetectorPercussive, DetectorSoft
        , PhaseLaminar, PhaseIndependent
        -- threading omitted
        , WindowStandard, WindowShort, WindowLong
        , SmoothingOff, SmoothingOn
        , FormantShifted, FormantPreserved
        , PitchHighSpeed, PitchHighQuality, PitchHighConsistency
        , ChannelsApart, ChannelsTogether
        )
    , percussiveOptions
    , offline
) where
import qualified Control.Exception as Exception
import qualified Control.Monad.Trans.Resource as Resource
import qualified Data.Vector.Storable as V
import qualified Foreign
import           Foreign (Ptr)
import           GHC.TypeLits (KnownNat)
import qualified Streaming.Prelude as S

import qualified Util.Audio.Audio as Audio
import qualified Util.Audio.RubberbandC as RubberbandC
import           Util.Audio.RubberbandC (percussiveOptions, Option(..))
import qualified Util.Control as Control
import qualified Util.Lists as Lists

import           Global


data Config = Config {
    Config -> [Option]
_options :: ![Option]
    , Config -> Double
_timeRatio :: !Double
    , Config -> Double
_pitchRatio :: !Double
    } deriving (Int -> Config -> ShowS
[Config] -> ShowS
Config -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Config] -> ShowS
$cshowList :: [Config] -> ShowS
show :: Config -> String
$cshow :: Config -> String
showsPrec :: Int -> Config -> ShowS
$cshowsPrec :: Int -> Config -> ShowS
Show)

type Samples = V.Vector Audio.Sample

config :: Config
config :: Config
config = [Option] -> Double -> Double -> Config
Config [] Double
1 Double
1

{- Options:

    StretchElastic - offline only, variable stretch to preserve transients
    StretchPrecise - online mode, try to get linear stretch rate throughout

    TransientsCrisp - good for percussive
    TransientsMixed
    TransientsSmooth - good for melodic

    DetectorCompound - general transient detector, default
    DetectorPercussive
    DetectorSoft - for melodic

    ThreadingAuto - one thread per channel in offline
    ThreadingNever
    ThreadingAlways

    WindowStandard
    WindowShort - better for timing
    WindowLong

    FormantShifted - default, no formant processing
    FormantPreserved - maybe better for pitch shifting instruments?

    PitchHigh{Speed,Quality,Consistency} - streaming only
-}


-- TODO
-- take key frame map as a signal?

-- | Stretch or pitch shift audio in a non-streaming way.  This collects
-- the whole stream, does the transformation, and then streams the output.
offline :: forall rate chan. (KnownNat rate, KnownNat chan)
    => Config -> Audio.AudioIO rate chan -> Audio.AudioIO rate chan
offline :: forall (rate :: Nat) (chan :: Nat).
(KnownNat rate, KnownNat chan) =>
Config -> AudioIO rate chan -> AudioIO rate chan
offline Config
config AudioIO rate chan
audio = forall (m :: * -> *) (rate :: Nat) (chan :: Nat).
Stream (Of Block) m () -> Audio m rate chan
Audio.Audio forall a b. (a -> b) -> a -> b
$ do
    [Block]
blocks <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) (rate :: Nat) (chan :: Nat).
Monad m =>
Audio m rate chan -> m [Block]
Audio.toBlocks AudioIO rate chan
audio
    forall (m :: * -> *).
Monad m =>
Stream (Of [Samples]) m () -> Stream (Of Block) m ()
interleave forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> Config
-> [Samples]
-> Stream (Of [Samples]) (ResourceT IO) ()
offlineStream Int
rate Int
chan Config
config (Int -> [Block] -> [Samples]
deinterleave Int
chan [Block]
blocks)
    where
    rate :: Int
rate = forall (n :: Nat). KnownNat n => Proxy n -> Int
Audio.natVal (forall {k} (t :: k). Proxy t
Proxy @rate)
    chan :: Int
chan = forall (n :: Nat). KnownNat n => Proxy n -> Int
Audio.natVal (forall {k} (t :: k). Proxy t
Proxy @chan)

interleave :: Monad m => S.Stream (S.Of [Samples]) m ()
    -> S.Stream (S.Of Audio.Block) m ()
interleave :: forall (m :: * -> *).
Monad m =>
Stream (Of [Samples]) m () -> Stream (Of Block) m ()
interleave = forall (m :: * -> *) a b r.
Monad m =>
(a -> b) -> Stream (Of a) m r -> Stream (Of b) m r
S.map (Samples -> Block
Audio.Block forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Storable a => [Vector a] -> Vector a
Audio.interleaveV)

-- | This reallocates the whole audio stream as one big vector for each
-- channel.  I could avoid this by passing it into rubberband_study and
-- rubberband_process as chunks, but since they all have to be kept in memory
-- anyway, I don't think it would help, and it's more complicated to implement.
deinterleave :: Audio.Channels -> [Audio.Block] -> [Samples]
deinterleave :: Int -> [Block] -> [Samples]
deinterleave Int
chan = forall a b. (a -> b) -> [a] -> [b]
map (forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map Block -> Samples
Audio.blockVector) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [[a]] -> [[a]]
Lists.rotate
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (Int -> Block -> [Block]
Audio.deinterleaveB Int
chan)

offlineStream :: Audio.Rate -> Audio.Channels -> Config -> [Samples]
    -> S.Stream (S.Of [Samples]) (Resource.ResourceT IO) ()
offlineStream :: Int
-> Int
-> Config
-> [Samples]
-> Stream (Of [Samples]) (ResourceT IO) ()
offlineStream Int
_ Int
_ Config
_ [] = forall (m :: * -> *) a. Monad m => a -> m a
return ()
offlineStream Int
srate Int
chan Config
config inputs :: [Samples]
inputs@(Samples
input0:[Samples]
_) = do
    (ReleaseKey
key, State
state) <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadResource m =>
IO a -> (a -> IO ()) -> m (ReleaseKey, a)
Resource.allocate IO State
new State -> IO ()
RubberbandC.rubberband_delete
    -- Debug.traceM "set_expected" totalFrames
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ State -> Frames -> IO ()
RubberbandC.rubberband_set_expected_input_duration State
state
        Frames
totalFrames
    -- Debug.traceM "study" ()
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall val a.
Storable val =>
[Vector val] -> Int -> (Ptr (Ptr val) -> IO a) -> IO a
incrementPtrs [Samples]
inputs Int
0 forall a b. (a -> b) -> a -> b
$ \Ptr (Ptr Sample)
inputpp ->
        State -> Ptr (Ptr Sample) -> Frames -> Bool -> IO ()
RubberbandC.rubberband_study State
state Ptr (Ptr Sample)
inputpp Frames
totalFrames Bool
True
    forall state a. state -> ((state -> a) -> state -> a) -> a
Control.loop1 Frames
0 (forall {m :: * -> *}.
MonadIO m =>
ReleaseKey
-> State
-> [Samples]
-> (Frames -> Stream (Of [Samples]) m ())
-> Frames
-> Stream (Of [Samples]) m ()
process ReleaseKey
key State
state [Samples]
inputs)
    where
    process :: ReleaseKey
-> State
-> [Samples]
-> (Frames -> Stream (Of [Samples]) m ())
-> Frames
-> Stream (Of [Samples]) m ()
process ReleaseKey
key State
state [Samples]
inputs Frames -> Stream (Of [Samples]) m ()
loop Frames
now = do
        Frames
required <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ State -> IO Frames
RubberbandC.rubberband_get_samples_required State
state
        let left :: Frames
left = Frames
totalFrames forall a. Num a => a -> a -> a
- Frames
now
        -- Debug.traceM "required" (now, required, left)
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Frames
required forall a. Ord a => a -> a -> Bool
> Frames
0 Bool -> Bool -> Bool
&& Frames
left forall a. Ord a => a -> a -> Bool
> Frames
0) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
            forall val a.
Storable val =>
[Vector val] -> Int -> (Ptr (Ptr val) -> IO a) -> IO a
incrementPtrs [Samples]
inputs (forall a b. (Integral a, Num b) => a -> b
fromIntegral Frames
now) forall a b. (a -> b) -> a -> b
$ \Ptr (Ptr Sample)
inputpp ->
                State -> Ptr (Ptr Sample) -> Frames -> Bool -> IO ()
RubberbandC.rubberband_process State
state Ptr (Ptr Sample)
inputpp
                    (forall a. Ord a => a -> a -> a
min Frames
required Frames
left)
                    (Frames
now forall a. Num a => a -> a -> a
+ Frames
required forall a. Ord a => a -> a -> Bool
>= Frames
totalFrames)
            -- Debug.traceM "process"
            --     (now, min required left, now + required >= totalFrames)
        Int
available <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ State -> IO Int
RubberbandC.rubberband_available State
state
        -- Debug.traceM "available" available
        if Int
available forall a. Eq a => a -> a -> Bool
== -Int
1
            then forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (forall (m :: * -> *). MonadIO m => ReleaseKey -> m ()
Resource.release ReleaseKey
key)
            else do
                forall (m :: * -> *) a. Monad m => a -> Stream (Of a) m ()
S.yield forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (State -> Int -> IO [Samples]
retrieve State
state Int
available)
                Frames -> Stream (Of [Samples]) m ()
loop forall a b. (a -> b) -> a -> b
$ Frames
now forall a. Num a => a -> a -> a
+ forall a. Ord a => a -> a -> a
min Frames
required Frames
left
    retrieve :: State -> Int -> IO [Samples]
retrieve State
state Int
available = do
        [Ptr Sample]
outputsps <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
chan (forall a. Storable a => Int -> IO (Ptr a)
Foreign.mallocArray Int
available)
        forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
Foreign.withArray [Ptr Sample]
outputsps forall a b. (a -> b) -> a -> b
$ \Ptr (Ptr Sample)
outputspp -> do
            Frames
retrieved <- State -> Ptr (Ptr Sample) -> Frames -> IO Frames
RubberbandC.rubberband_retrieve State
state Ptr (Ptr Sample)
outputspp
                (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
available)
            [ForeignPtr Sample]
outFPs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
Foreign.newForeignPtr forall a. FinalizerPtr a
Foreign.finalizerFree)
                [Ptr Sample]
outputsps
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map
                (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. ForeignPtr a -> Int -> Vector a
V.unsafeFromForeignPtr0 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Frames
retrieved))
                [ForeignPtr Sample]
outFPs

    totalFrames :: Frames
totalFrames = forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Storable a => Vector a -> Int
V.length Samples
input0)
    new :: IO State
new = Frames -> Int -> [Option] -> Double -> Double -> IO State
RubberbandC.rubberband_new (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
srate) Int
chan
        (Option
ProcessOffline forall a. a -> [a] -> [a]
: Config -> [Option]
_options Config
config) (Config -> Double
_timeRatio Config
config)
        (Config -> Double
_pitchRatio Config
config)

-- | Convert vectors to ptrs, incremented by some index.
incrementPtrs :: Foreign.Storable val => [V.Vector val] -> Int
    -> (Ptr (Ptr val) -> IO a) -> IO a
incrementPtrs :: forall val a.
Storable val =>
[Vector val] -> Int -> (Ptr (Ptr val) -> IO a) -> IO a
incrementPtrs [Vector val]
vs Int
increment Ptr (Ptr val) -> IO a
action = forall a b res.
(a -> (b -> res) -> res) -> [a] -> ([b] -> res) -> res
Foreign.withMany forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
V.unsafeWith [Vector val]
vs forall a b. (a -> b) -> a -> b
$ \[Ptr val]
vps ->
    forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
Foreign.withArray (forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Storable a => Ptr a -> Int -> Ptr a
Foreign.advancePtr Int
increment) [Ptr val]
vps) Ptr (Ptr val) -> IO a
action


-- * old

-- I have to collect the whole input, but then I can stream the output.
-- So I could output as a callback, or I could just return a stream.
offlineList :: Audio.Rate -> Config -> [Samples] -> IO [[Samples]]
offlineList :: Int -> Config -> [Samples] -> IO [[Samples]]
offlineList Int
_ Config
_ [] = forall (m :: * -> *) a. Monad m => a -> m a
return []
offlineList Int
srate Config
config [Samples]
inputs =
    forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
Exception.bracket IO State
new State -> IO ()
RubberbandC.rubberband_delete forall a b. (a -> b) -> a -> b
$ \State
state -> do
        State -> Frames -> IO ()
RubberbandC.rubberband_set_expected_input_duration State
state Frames
totalFrames
        forall a b res.
(a -> (b -> res) -> res) -> [a] -> ([b] -> res) -> res
Foreign.withMany forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
V.unsafeWith [Samples]
inputs forall a b. (a -> b) -> a -> b
$ \[Ptr Sample]
inputps -> do
            forall val a.
Storable val =>
[Ptr val] -> Int -> (Ptr (Ptr val) -> IO a) -> IO a
incrementPtrs2 [Ptr Sample]
inputps Int
0 forall a b. (a -> b) -> a -> b
$ \Ptr (Ptr Sample)
inputpp ->
                State -> Ptr (Ptr Sample) -> Frames -> Bool -> IO ()
RubberbandC.rubberband_study State
state Ptr (Ptr Sample)
inputpp Frames
totalFrames Bool
True
            forall s1 s2 a. s1 -> s2 -> ((s1 -> s2 -> a) -> s1 -> s2 -> a) -> a
Control.loop2 [] Frames
0 (State
-> [Ptr Sample]
-> ([[Samples]] -> Frames -> IO [[Samples]])
-> [[Samples]]
-> Frames
-> IO [[Samples]]
process State
state [Ptr Sample]
inputps)
    where
    process :: State
-> [Ptr Sample]
-> ([[Samples]] -> Frames -> IO [[Samples]])
-> [[Samples]]
-> Frames
-> IO [[Samples]]
process State
state [Ptr Sample]
inputps [[Samples]] -> Frames -> IO [[Samples]]
loop [[Samples]]
accum Frames
now = do
        Frames
wanted <- State -> IO Frames
RubberbandC.rubberband_get_samples_required State
state
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Frames
wanted forall a. Eq a => a -> a -> Bool
== Frames
0) forall a b. (a -> b) -> a -> b
$
            forall val a.
Storable val =>
[Ptr val] -> Int -> (Ptr (Ptr val) -> IO a) -> IO a
incrementPtrs2 [Ptr Sample]
inputps (forall a b. (Integral a, Num b) => a -> b
fromIntegral Frames
now) forall a b. (a -> b) -> a -> b
$ \Ptr (Ptr Sample)
inputpp ->
                State -> Ptr (Ptr Sample) -> Frames -> Bool -> IO ()
RubberbandC.rubberband_process State
state Ptr (Ptr Sample)
inputpp Frames
wanted
                    (Frames
now forall a. Num a => a -> a -> a
+ Frames
wanted forall a. Ord a => a -> a -> Bool
>= Frames
totalFrames)
        Int
available <- State -> IO Int
RubberbandC.rubberband_available State
state
        if Int
available forall a. Eq a => a -> a -> Bool
== -Int
1 then forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. [a] -> [a]
reverse [[Samples]]
accum) else do
            [Ptr Sample]
outputsps <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
chan (forall a. Storable a => Int -> IO (Ptr a)
Foreign.mallocArray Int
available)
            ([Samples]
outputs, Frames
retrieved)<-forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
Foreign.withArray [Ptr Sample]
outputsps forall a b. (a -> b) -> a -> b
$ \Ptr (Ptr Sample)
outputspp -> do
                Frames
retrieved <- State -> Ptr (Ptr Sample) -> Frames -> IO Frames
RubberbandC.rubberband_retrieve State
state Ptr (Ptr Sample)
outputspp
                    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
available)
                [ForeignPtr Sample]
outFPs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
Foreign.newForeignPtr forall a. FinalizerPtr a
Foreign.finalizerFree)
                    [Ptr Sample]
outputsps
                let outputs :: [Samples]
outputs = forall a b. (a -> b) -> [a] -> [b]
map
                        (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. ForeignPtr a -> Int -> Vector a
V.unsafeFromForeignPtr0 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Frames
retrieved))
                        [ForeignPtr Sample]
outFPs
                forall (m :: * -> *) a. Monad m => a -> m a
return ([Samples]
outputs, Frames
retrieved)
            [[Samples]] -> Frames -> IO [[Samples]]
loop ([Samples]
outputsforall a. a -> [a] -> [a]
:[[Samples]]
accum) (Frames
now forall a. Num a => a -> a -> a
+ Frames
retrieved)

    totalFrames :: Frames
totalFrames = forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Storable a => Vector a -> Int
V.length (forall a. [a] -> a
head [Samples]
inputs))
    new :: IO State
new = Frames -> Int -> [Option] -> Double -> Double -> IO State
RubberbandC.rubberband_new (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
srate) Int
chan
        (Option
ProcessOffline forall a. a -> [a] -> [a]
: Config -> [Option]
_options Config
config) (Config -> Double
_timeRatio Config
config)
        (Config -> Double
_pitchRatio Config
config)
    chan :: Int
chan = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Samples]
inputs

incrementPtrs2 :: Foreign.Storable val => [Ptr val] -> Int
    -> (Ptr (Ptr val) -> IO a) -> IO a
incrementPtrs2 :: forall val a.
Storable val =>
[Ptr val] -> Int -> (Ptr (Ptr val) -> IO a) -> IO a
incrementPtrs2 [Ptr val]
vps Int
increment Ptr (Ptr val) -> IO a
action = do
    forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
Foreign.withArray (forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Storable a => Ptr a -> Int -> Ptr a
Foreign.advancePtr Int
increment) [Ptr val]
vps) Ptr (Ptr val) -> IO a
action