{-# LANGUAGE LambdaCase #-}

module Language.Bitcoin.Script.Descriptors.Utils (
    descriptorAddresses,
    compile,
) where

import Data.List (sortOn)
import Data.Maybe (maybeToList)
import Data.Serialize (decode, encode)
import Haskoin (
    Address,
    Script,
    ScriptOutput (..),
    addressHash,
    eitherToMaybe,
    payToNestedScriptAddress,
    payToScriptAddress,
    payToWitnessScriptAddress,
    pubKeyAddr,
    pubKeyCompatWitnessAddr,
    pubKeyCompressed,
    pubKeyWitnessAddr,
    sortMulSig,
 )

import qualified Language.Bitcoin.Miniscript as M
import Language.Bitcoin.Script.Descriptors.Syntax (
    OutputDescriptor (..),
    ScriptDescriptor (..),
    keyBytes,
    keyDescPubKey,
 )

{- | Get the set of addresses associated with an output descriptor.  The list will be empty if:

     * any keys are indefinite
     * the output is p2pk
     * the output has a non-standard script

     The list can contain more than one address in the case of the "combo" construct.
-}
descriptorAddresses :: OutputDescriptor -> [Address]
descriptorAddresses :: OutputDescriptor -> [Address]
descriptorAddresses = \case
    ScriptPubKey Pk{} -> [Address]
forall a. Monoid a => a
mempty
    ScriptPubKey (Pkh KeyDescriptor
key) -> (PubKeyI -> [Address]) -> Maybe PubKeyI -> [Address]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Address -> [Address]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Address -> [Address])
-> (PubKeyI -> Address) -> PubKeyI -> [Address]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PubKeyI -> Address
pubKeyAddr) (Maybe PubKeyI -> [Address]) -> Maybe PubKeyI -> [Address]
forall a b. (a -> b) -> a -> b
$ KeyDescriptor -> Maybe PubKeyI
keyDescPubKey KeyDescriptor
key
    P2SH ScriptDescriptor
descriptor -> Maybe Address -> [Address]
forall a. Maybe a -> [a]
maybeToList (Maybe Address -> [Address]) -> Maybe Address -> [Address]
forall a b. (a -> b) -> a -> b
$ ScriptOutput -> Address
payToScriptAddress (ScriptOutput -> Address) -> Maybe ScriptOutput -> Maybe Address
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ScriptDescriptor -> Maybe ScriptOutput
scriptDescriptorOutput ScriptDescriptor
descriptor
    P2WPKH KeyDescriptor
key -> (PubKeyI -> [Address]) -> Maybe PubKeyI -> [Address]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Address -> [Address]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Address -> [Address])
-> (PubKeyI -> Address) -> PubKeyI -> [Address]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PubKeyI -> Address
pubKeyWitnessAddr) (Maybe PubKeyI -> [Address]) -> Maybe PubKeyI -> [Address]
forall a b. (a -> b) -> a -> b
$ KeyDescriptor -> Maybe PubKeyI
keyDescPubKey KeyDescriptor
key
    P2WSH ScriptDescriptor
descriptor -> Maybe Address -> [Address]
forall a. Maybe a -> [a]
maybeToList (Maybe Address -> [Address]) -> Maybe Address -> [Address]
forall a b. (a -> b) -> a -> b
$ ScriptOutput -> Address
payToWitnessScriptAddress (ScriptOutput -> Address) -> Maybe ScriptOutput -> Maybe Address
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ScriptDescriptor -> Maybe ScriptOutput
scriptDescriptorOutput ScriptDescriptor
descriptor
    WrappedWPkh KeyDescriptor
key -> (PubKeyI -> [Address]) -> Maybe PubKeyI -> [Address]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Address -> [Address]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Address -> [Address])
-> (PubKeyI -> Address) -> PubKeyI -> [Address]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PubKeyI -> Address
pubKeyCompatWitnessAddr) (Maybe PubKeyI -> [Address]) -> Maybe PubKeyI -> [Address]
forall a b. (a -> b) -> a -> b
$ KeyDescriptor -> Maybe PubKeyI
keyDescPubKey KeyDescriptor
key
    WrappedWSh ScriptDescriptor
descriptor -> Maybe Address -> [Address]
forall a. Maybe a -> [a]
maybeToList (Maybe Address -> [Address]) -> Maybe Address -> [Address]
forall a b. (a -> b) -> a -> b
$ ScriptOutput -> Address
payToNestedScriptAddress (ScriptOutput -> Address) -> Maybe ScriptOutput -> Maybe Address
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ScriptDescriptor -> Maybe ScriptOutput
scriptDescriptorOutput ScriptDescriptor
descriptor
    Combo KeyDescriptor
key
        | Just PubKeyI
pk <- KeyDescriptor -> Maybe PubKeyI
keyDescPubKey KeyDescriptor
key ->
            [PubKeyI -> Address
pubKeyAddr PubKeyI
pk]
                [Address] -> [Address] -> [Address]
forall a. Semigroup a => a -> a -> a
<> if PubKeyI -> Bool
pubKeyCompressed PubKeyI
pk
                    then [PubKeyI -> Address
pubKeyWitnessAddr PubKeyI
pk, PubKeyI -> Address
pubKeyCompatWitnessAddr PubKeyI
pk]
                    else [Address]
forall a. Monoid a => a
mempty
    Addr Address
addr -> [Address
addr]
    OutputDescriptor
_ -> [Address]
forall a. Monoid a => a
mempty

scriptDescriptorOutput :: ScriptDescriptor -> Maybe ScriptOutput
scriptDescriptorOutput :: ScriptDescriptor -> Maybe ScriptOutput
scriptDescriptorOutput = \case
    Pk KeyDescriptor
key -> PubKeyI -> ScriptOutput
PayPK (PubKeyI -> ScriptOutput) -> Maybe PubKeyI -> Maybe ScriptOutput
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> KeyDescriptor -> Maybe PubKeyI
keyDescPubKey KeyDescriptor
key
    Pkh KeyDescriptor
key -> Hash160 -> ScriptOutput
PayPKHash (Hash160 -> ScriptOutput)
-> (PubKeyI -> Hash160) -> PubKeyI -> ScriptOutput
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Hash160
forall b. ByteArrayAccess b => b -> Hash160
addressHash (ByteString -> Hash160)
-> (PubKeyI -> ByteString) -> PubKeyI -> Hash160
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PubKeyI -> ByteString
forall a. Serialize a => a -> ByteString
encode (PubKeyI -> ScriptOutput) -> Maybe PubKeyI -> Maybe ScriptOutput
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> KeyDescriptor -> Maybe PubKeyI
keyDescPubKey KeyDescriptor
key
    Multi Int
k [KeyDescriptor]
ks -> [PubKeyI] -> Int -> ScriptOutput
PayMulSig ([PubKeyI] -> Int -> ScriptOutput)
-> Maybe [PubKeyI] -> Maybe (Int -> ScriptOutput)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (KeyDescriptor -> Maybe PubKeyI)
-> [KeyDescriptor] -> Maybe [PubKeyI]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse KeyDescriptor -> Maybe PubKeyI
keyDescPubKey [KeyDescriptor]
ks Maybe (Int -> ScriptOutput) -> Maybe Int -> Maybe ScriptOutput
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> Maybe Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
k
    SortedMulti Int
k [KeyDescriptor]
ks -> ScriptOutput -> ScriptOutput
sortMulSig (ScriptOutput -> ScriptOutput)
-> Maybe ScriptOutput -> Maybe ScriptOutput
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([PubKeyI] -> Int -> ScriptOutput
PayMulSig ([PubKeyI] -> Int -> ScriptOutput)
-> Maybe [PubKeyI] -> Maybe (Int -> ScriptOutput)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (KeyDescriptor -> Maybe PubKeyI)
-> [KeyDescriptor] -> Maybe [PubKeyI]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse KeyDescriptor -> Maybe PubKeyI
keyDescPubKey [KeyDescriptor]
ks Maybe (Int -> ScriptOutput) -> Maybe Int -> Maybe ScriptOutput
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> Maybe Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
k)
    ScriptDescriptor
_ -> Maybe ScriptOutput
forall a. Maybe a
Nothing

-- | Produce the script described by the descriptor.  Fails when any keys in the descriptor are indeterminate.
compile :: ScriptDescriptor -> Maybe Script
compile :: ScriptDescriptor -> Maybe Script
compile = \case
    Pk KeyDescriptor
key -> Miniscript -> Maybe Script
compileMaybe (Miniscript -> Maybe Script) -> Miniscript -> Maybe Script
forall a b. (a -> b) -> a -> b
$ KeyDescriptor -> Miniscript
M.key KeyDescriptor
key
    Pkh KeyDescriptor
key -> Miniscript -> Maybe Script
compileMaybe (Miniscript -> Maybe Script) -> Miniscript -> Maybe Script
forall a b. (a -> b) -> a -> b
$ KeyDescriptor -> Miniscript
M.keyH KeyDescriptor
key
    Multi Int
k [KeyDescriptor]
ks -> Miniscript -> Maybe Script
compileMaybe (Miniscript -> Maybe Script) -> Miniscript -> Maybe Script
forall a b. (a -> b) -> a -> b
$ Int -> [KeyDescriptor] -> Miniscript
M.multi Int
k [KeyDescriptor]
ks
    SortedMulti Int
k [KeyDescriptor]
ks -> Miniscript -> Maybe Script
compileMaybe (Miniscript -> Maybe Script) -> Miniscript -> Maybe Script
forall a b. (a -> b) -> a -> b
$ Int -> [KeyDescriptor] -> Miniscript
M.multi Int
k ((KeyDescriptor -> Maybe ByteString)
-> [KeyDescriptor] -> [KeyDescriptor]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn KeyDescriptor -> Maybe ByteString
keyBytes [KeyDescriptor]
ks)
    Raw ByteString
bs -> Either String Script -> Maybe Script
forall a b. Either a b -> Maybe b
eitherToMaybe (ByteString -> Either String Script
forall a. Serialize a => ByteString -> Either String a
decode ByteString
bs)
  where
    compileMaybe :: Miniscript -> Maybe Script
compileMaybe = Either CompilerError Script -> Maybe Script
forall a b. Either a b -> Maybe b
eitherToMaybe (Either CompilerError Script -> Maybe Script)
-> (Miniscript -> Either CompilerError Script)
-> Miniscript
-> Maybe Script
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Miniscript -> Either CompilerError Script
M.compile