<commit_before>package itest

import (
	"context"
	"math"

	"github.com/btcsuite/btcd/btcutil"
	"github.com/btcsuite/btcd/btcutil/hdkeychain"
	"github.com/lightningnetwork/lnd/aezeed"
	"github.com/lightningnetwork/lnd/lnrpc"
	"github.com/lightningnetwork/lnd/lntest"
	"github.com/lightningnetwork/lnd/lntest/wait"
	"github.com/stretchr/testify/require"
)

// testGetRecoveryInfo checks whether lnd gives the right information about
// the wallet recovery process.
func testGetRecoveryInfo(net *lntest.NetworkHarness, t *harnessTest) {
	ctxb := context.Background()

	// First, create a new node with strong passphrase and grab the mnemonic
	// used for key derivation. This will bring up Carol with an empty
	// wallet, and such that she is synced up.
	password := []byte("The Magic Words are Squeamish Ossifrage")
	carol, mnemonic, _, err := net.NewNodeWithSeed(
		"Carol", nil, password, false,
	)
	if err != nil {
		t.Fatalf("unable to create node with seed; %v", err)
	}

	shutdownAndAssert(net, t, carol)

	checkInfo := func(expectedRecoveryMode, expectedRecoveryFinished bool,
		expectedProgress float64, recoveryWindow int32) {

		// Restore Carol, passing in the password, mnemonic, and
		// desired recovery window.
		node, err := net.RestoreNodeWithSeed(
			"Carol", nil, password, mnemonic, "", recoveryWindow,
			nil,
		)
		if err != nil {
			t.Fatalf("unable to restore node: %v", err)
		}

		// Wait for Carol to sync to the chain.
		_, minerHeight, err := net.Miner.Client.GetBestBlock()
		if err != nil {
			t.Fatalf("unable to get current blockheight %v", err)
		}
		err = waitForNodeBlockHeight(node, minerHeight)
		if err != nil {
			t.Fatalf("unable to sync to chain: %v", err)
		}

		// Query carol for her current wallet recovery progress.
		var (
			recoveryMode     bool
			recoveryFinished bool
			progress         float64
		)

		err = wait.Predicate(func() bool {
			// Verify that recovery info gives the right response.
			req := &lnrpc.GetRecoveryInfoRequest{}
			ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
			resp, err := node.GetRecoveryInfo(ctxt, req)
			if err != nil {
				t.Fatalf("unable to query recovery info: %v", err)
			}

			recoveryMode = resp.RecoveryMode
			recoveryFinished = resp.RecoveryFinished
			progress = resp.Progress

			if recoveryMode != expectedRecoveryMode ||
				recoveryFinished != expectedRecoveryFinished ||
				progress != expectedProgress {

				return false
			}

			return true
		}, defaultTimeout)
		if err != nil {
			t.Fatalf("expected recovery mode to be %v, got %v, "+
				"expected recovery finished to be %v, got %v, "+
				"expected progress %v, got %v",
				expectedRecoveryMode, recoveryMode,
				expectedRecoveryFinished, recoveryFinished,
				expectedProgress, progress,
			)
		}

		// Lastly, shutdown this Carol so we can move on to the next
		// restoration.
		shutdownAndAssert(net, t, node)
	}

	// Restore Carol with a recovery window of 0. Since it's not in recovery
	// mode, the recovery info will give a response with recoveryMode=false,
	// recoveryFinished=false, and progress=0
	checkInfo(false, false, 0, 0)

	// Change the recovery windown to be 1 to turn on recovery mode. Since the
	// current chain height is the same as the birthday height, it should
	// indicate the recovery process is finished.
	checkInfo(true, true, 1, 1)

	// We now go ahead 5 blocks. Because the wallet's syncing process is
	// controlled by a goroutine in the background, it will catch up quickly.
	// This makes the recovery progress back to 1.
	mineBlocks(t, net, 5, 0)
	checkInfo(true, true, 1, 1)
}

// testOnchainFundRecovery checks lnd's ability to rescan for onchain outputs
// when providing a valid aezeed that owns outputs on the chain. This test
// performs multiple restorations using the same seed and various recovery
// windows to ensure we detect funds properly.
func testOnchainFundRecovery(net *lntest.NetworkHarness, t *harnessTest) {
	ctxb := context.Background()

	// First, create a new node with strong passphrase and grab the mnemonic
	// used for key derivation. This will bring up Carol with an empty
	// wallet, and such that she is synced up.
	password := []byte("The Magic Words are Squeamish Ossifrage")
	carol, mnemonic, _, err := net.NewNodeWithSeed(
		"Carol", nil, password, false,
	)
	require.NoError(t.t, err)
	shutdownAndAssert(net, t, carol)

	// As long as the mnemonic is non-nil and the extended key is empty, the
	// closure below will always restore the node from the seed. The tests
	// need to manually overwrite this value to change that behavior.
	rootKey := ""

	// Create a closure for testing the recovery of Carol's wallet. This
	// method takes the expected value of Carol's balance when using the
	// given recovery window. Additionally, the caller can specify an action
	// to perform on the restored node before the node is shutdown.
	restoreCheckBalance := func(expAmount int64, expectedNumUTXOs uint32,
		recoveryWindow int32, fn func(*lntest.HarnessNode)) {

		t.t.Helper()

		// Restore Carol, passing in the password, mnemonic, and
		// desired recovery window.
		node, err := net.RestoreNodeWithSeed(
			"Carol", nil, password, mnemonic, rootKey,
			recoveryWindow, nil,
		)
		require.NoError(t.t, err)

		// Query carol for her current wallet balance, and also that we
		// gain the expected number of UTXOs.
		var (
			currBalance  int64
			currNumUTXOs uint32
		)
		err = wait.Predicate(func() bool {
			req := &lnrpc.WalletBalanceRequest{}
			ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
			resp, err := node.WalletBalance(ctxt, req)
			require.NoError(t.t, err)
			currBalance = resp.ConfirmedBalance

			utxoReq := &lnrpc.ListUnspentRequest{
				MaxConfs: math.MaxInt32,
			}
			ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
			utxoResp, err := node.ListUnspent(ctxt, utxoReq)
			require.NoError(t.t, err)
			currNumUTXOs = uint32(len(utxoResp.Utxos))

			// Verify that Carol's balance and number of UTXOs
			// matches what's expected.
			if expAmount != currBalance {
				return false
			}
			if currNumUTXOs != expectedNumUTXOs {
				return false
			}

			return true
		}, defaultTimeout)
		if err != nil {
			t.Fatalf("expected restored node to have %d satoshis, "+
				"instead has %d satoshis, expected %d utxos "+
				"instead has %d", expAmount, currBalance,
				expectedNumUTXOs, currNumUTXOs)
		}

		// If the user provided a callback, execute the commands against
		// the restored Carol.
		if fn != nil {
			fn(node)
		}

		// Lastly, shutdown this Carol so we can move on to the next
		// restoration.
		shutdownAndAssert(net, t, node)
	}

	// Create a closure-factory for building closures that can generate and
	// skip a configurable number of addresses, before finally sending coins
	// to a next generated address. The returned closure will apply the same
	// behavior to both default P2WKH and NP2WKH scopes.
	skipAndSend := func(nskip int) func(*lntest.HarnessNode) {
		return func(node *lntest.HarnessNode) {
			t.t.Helper()

			newP2WKHAddrReq := &lnrpc.NewAddressRequest{
				Type: AddrTypeWitnessPubkeyHash,
			}

			newNP2WKHAddrReq := &lnrpc.NewAddressRequest{
				Type: AddrTypeNestedPubkeyHash,
			}

			newP2TRAddrReq := &lnrpc.NewAddressRequest{
				Type: AddrTypeTaprootPubkey,
			}

			// Generate and skip the number of addresses requested.
			ctxt, cancel := context.WithTimeout(
				ctxb, defaultTimeout,
			)
			defer cancel()
			for i := 0; i < nskip; i++ {
				_, err = node.NewAddress(ctxt, newP2WKHAddrReq)
				require.NoError(t.t, err)

				_, err = node.NewAddress(ctxt, newNP2WKHAddrReq)
				require.NoError(t.t, err)

				_, err = node.NewAddress(ctxt, newP2TRAddrReq)
				require.NoError(t.t, err)
			}

			// Send one BTC to the next P2WKH address.
			net.SendCoins(t.t, btcutil.SatoshiPerBitcoin, node)

			// And another to the next NP2WKH address.
			net.SendCoinsNP2WKH(
				t.t, btcutil.SatoshiPerBitcoin, node,
			)

			// Add another whole coin to the P2TR address.
			net.SendCoinsP2TR(
				t.t, btcutil.SatoshiPerBitcoin, node,
			)
		}
	}

	// Restore Carol with a recovery window of 0. Since no coins have been
	// sent, her balance should be zero.
	//
	// After, one BTC is sent to both her first external P2WKH and NP2WKH
	// addresses.
	restoreCheckBalance(0, 0, 0, skipAndSend(0))

	// Check that restoring without a look-ahead results in having no funds
	// in the wallet, even though they exist on-chain.
	restoreCheckBalance(0, 0, 0, nil)

	// Now, check that using a look-ahead of 1 recovers the balance from
	// the two transactions above. We should also now have 2 UTXOs in the
	// wallet at the end of the recovery attempt.
	//
	// After, we will generate and skip 9 P2WKH, NP2WKH and P2TR addresses,
	// and send another BTC to the subsequent 10th address in each
	// derivation path.
	restoreCheckBalance(3*btcutil.SatoshiPerBitcoin, 3, 1, skipAndSend(9))

	// Check that using a recovery window of 9 does not find the two most
	// recent txns.
	restoreCheckBalance(3*btcutil.SatoshiPerBitcoin, 3, 9, nil)

	// Extending our recovery window to 10 should find the most recent
	// transactions, leaving the wallet with 6 BTC total. We should also
	// learn of the two additional UTXOs created above.
	//
	// After, we will skip 19 more addrs, sending to the 20th address past
	// our last found address, and repeat the same checks.
	restoreCheckBalance(6*btcutil.SatoshiPerBitcoin, 6, 10, skipAndSend(19))

	// Check that recovering with a recovery window of 19 fails to find the
	// most recent transactions.
	restoreCheckBalance(6*btcutil.SatoshiPerBitcoin, 6, 19, nil)

	// Ensure that using a recovery window of 20 succeeds with all UTXOs
	// found and the final balance reflected.

	// After these checks are done, we'll want to make sure we can also
	// recover change address outputs.  This is mainly motivated by a now
	// fixed bug in the wallet in which change addresses could at times be
	// created outside of the default key scopes. Recovery only used to be
	// performed on the default key scopes, so ideally this test case
	// would've caught the bug earlier. Carol has received 9 BTC so far from
	// the miner, we'll send 8 back to ensure all of her UTXOs get spent to
	// avoid fee discrepancies and a change output is formed.
	const minerAmt = 8 * btcutil.SatoshiPerBitcoin
	const finalBalance = 9 * btcutil.SatoshiPerBitcoin
	promptChangeAddr := func(node *lntest.HarnessNode) {
		t.t.Helper()

		minerAddr, err := net.Miner.NewAddress()
		require.NoError(t.t, err)
		ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
		resp, err := node.SendCoins(ctxt, &lnrpc.SendCoinsRequest{
			Addr:   minerAddr.String(),
			Amount: minerAmt,
		})
		require.NoError(t.t, err)
		txid, err := waitForTxInMempool(
			net.Miner.Client, minerMempoolTimeout,
		)
		require.NoError(t.t, err)
		require.Equal(t.t, txid.String(), resp.Txid)

		block := mineBlocks(t, net, 1, 1)[0]
		assertTxInBlock(t, block, txid)
	}
	restoreCheckBalance(finalBalance, 9, 20, promptChangeAddr)

	// We should expect a static fee of 50100 satoshis for spending 9 inputs
	// (3 P2WPKH, 3 NP2WPKH, 3 P2TR) to two P2WPKH outputs. Carol should
	// therefore only have one UTXO present (the change output) of
	// 9 - 8 - fee BTC.
	const fee = 50100
	restoreCheckBalance(finalBalance-minerAmt-fee, 1, 21, nil)

	// Last of all, make sure we can also restore a node from the extended
	// master root key directly instead of the seed.
	var seedMnemonic aezeed.Mnemonic
	copy(seedMnemonic[:], mnemonic)
	cipherSeed, err := seedMnemonic.ToCipherSeed(password)
	require.NoError(t.t, err)
	extendedRootKey, err := hdkeychain.NewMaster(
		cipherSeed.Entropy[:], harnessNetParams,
	)
	require.NoError(t.t, err)
	rootKey = extendedRootKey.String()
	mnemonic = nil

	restoreCheckBalance(finalBalance-minerAmt-fee, 1, 21, nil)
}
<commit_msg>lntest: adds test for previous outpoints<commit_after>package itest

import (
	"context"
	"math"

	"github.com/btcsuite/btcd/btcutil"
	"github.com/btcsuite/btcd/btcutil/hdkeychain"
	"github.com/lightningnetwork/lnd/aezeed"
	"github.com/lightningnetwork/lnd/lnrpc"
	"github.com/lightningnetwork/lnd/lntest"
	"github.com/lightningnetwork/lnd/lntest/wait"
	"github.com/stretchr/testify/require"
)

// testGetRecoveryInfo checks whether lnd gives the right information about
// the wallet recovery process.
func testGetRecoveryInfo(net *lntest.NetworkHarness, t *harnessTest) {
	ctxb := context.Background()

	// First, create a new node with strong passphrase and grab the mnemonic
	// used for key derivation. This will bring up Carol with an empty
	// wallet, and such that she is synced up.
	password := []byte("The Magic Words are Squeamish Ossifrage")
	carol, mnemonic, _, err := net.NewNodeWithSeed(
		"Carol", nil, password, false,
	)
	if err != nil {
		t.Fatalf("unable to create node with seed; %v", err)
	}

	shutdownAndAssert(net, t, carol)

	checkInfo := func(expectedRecoveryMode, expectedRecoveryFinished bool,
		expectedProgress float64, recoveryWindow int32) {

		// Restore Carol, passing in the password, mnemonic, and
		// desired recovery window.
		node, err := net.RestoreNodeWithSeed(
			"Carol", nil, password, mnemonic, "", recoveryWindow,
			nil,
		)
		if err != nil {
			t.Fatalf("unable to restore node: %v", err)
		}

		// Wait for Carol to sync to the chain.
		_, minerHeight, err := net.Miner.Client.GetBestBlock()
		if err != nil {
			t.Fatalf("unable to get current blockheight %v", err)
		}
		err = waitForNodeBlockHeight(node, minerHeight)
		if err != nil {
			t.Fatalf("unable to sync to chain: %v", err)
		}

		// Query carol for her current wallet recovery progress.
		var (
			recoveryMode     bool
			recoveryFinished bool
			progress         float64
		)

		err = wait.Predicate(func() bool {
			// Verify that recovery info gives the right response.
			req := &lnrpc.GetRecoveryInfoRequest{}
			ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
			resp, err := node.GetRecoveryInfo(ctxt, req)
			if err != nil {
				t.Fatalf("unable to query recovery info: %v", err)
			}

			recoveryMode = resp.RecoveryMode
			recoveryFinished = resp.RecoveryFinished
			progress = resp.Progress

			if recoveryMode != expectedRecoveryMode ||
				recoveryFinished != expectedRecoveryFinished ||
				progress != expectedProgress {

				return false
			}

			return true
		}, defaultTimeout)
		if err != nil {
			t.Fatalf("expected recovery mode to be %v, got %v, "+
				"expected recovery finished to be %v, got %v, "+
				"expected progress %v, got %v",
				expectedRecoveryMode, recoveryMode,
				expectedRecoveryFinished, recoveryFinished,
				expectedProgress, progress,
			)
		}

		// Lastly, shutdown this Carol so we can move on to the next
		// restoration.
		shutdownAndAssert(net, t, node)
	}

	// Restore Carol with a recovery window of 0. Since it's not in recovery
	// mode, the recovery info will give a response with recoveryMode=false,
	// recoveryFinished=false, and progress=0
	checkInfo(false, false, 0, 0)

	// Change the recovery windown to be 1 to turn on recovery mode. Since the
	// current chain height is the same as the birthday height, it should
	// indicate the recovery process is finished.
	checkInfo(true, true, 1, 1)

	// We now go ahead 5 blocks. Because the wallet's syncing process is
	// controlled by a goroutine in the background, it will catch up quickly.
	// This makes the recovery progress back to 1.
	mineBlocks(t, net, 5, 0)
	checkInfo(true, true, 1, 1)
}

// testOnchainFundRecovery checks lnd's ability to rescan for onchain outputs
// when providing a valid aezeed that owns outputs on the chain. This test
// performs multiple restorations using the same seed and various recovery
// windows to ensure we detect funds properly.
func testOnchainFundRecovery(net *lntest.NetworkHarness, t *harnessTest) {
	ctxb := context.Background()

	// First, create a new node with strong passphrase and grab the mnemonic
	// used for key derivation. This will bring up Carol with an empty
	// wallet, and such that she is synced up.
	password := []byte("The Magic Words are Squeamish Ossifrage")
	carol, mnemonic, _, err := net.NewNodeWithSeed(
		"Carol", nil, password, false,
	)
	require.NoError(t.t, err)
	shutdownAndAssert(net, t, carol)

	// As long as the mnemonic is non-nil and the extended key is empty, the
	// closure below will always restore the node from the seed. The tests
	// need to manually overwrite this value to change that behavior.
	rootKey := ""

	// Create a closure for testing the recovery of Carol's wallet. This
	// method takes the expected value of Carol's balance when using the
	// given recovery window. Additionally, the caller can specify an action
	// to perform on the restored node before the node is shutdown.
	restoreCheckBalance := func(expAmount int64, expectedNumUTXOs uint32,
		recoveryWindow int32, fn func(*lntest.HarnessNode)) {

		t.t.Helper()

		// Restore Carol, passing in the password, mnemonic, and
		// desired recovery window.
		node, err := net.RestoreNodeWithSeed(
			"Carol", nil, password, mnemonic, rootKey,
			recoveryWindow, nil,
		)
		require.NoError(t.t, err)

		// Query carol for her current wallet balance, and also that we
		// gain the expected number of UTXOs.
		var (
			currBalance  int64
			currNumUTXOs uint32
		)
		err = wait.Predicate(func() bool {
			req := &lnrpc.WalletBalanceRequest{}
			ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
			resp, err := node.WalletBalance(ctxt, req)
			require.NoError(t.t, err)
			currBalance = resp.ConfirmedBalance

			utxoReq := &lnrpc.ListUnspentRequest{
				MaxConfs: math.MaxInt32,
			}
			ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
			utxoResp, err := node.ListUnspent(ctxt, utxoReq)
			require.NoError(t.t, err)
			currNumUTXOs = uint32(len(utxoResp.Utxos))

			// Verify that Carol's balance and number of UTXOs
			// matches what's expected.
			if expAmount != currBalance {
				return false
			}
			if currNumUTXOs != expectedNumUTXOs {
				return false
			}

			return true
		}, defaultTimeout)
		if err != nil {
			t.Fatalf("expected restored node to have %d satoshis, "+
				"instead has %d satoshis, expected %d utxos "+
				"instead has %d", expAmount, currBalance,
				expectedNumUTXOs, currNumUTXOs)
		}

		// If the user provided a callback, execute the commands against
		// the restored Carol.
		if fn != nil {
			fn(node)
		}

		// Check if the previous outpoints are set correctly.
		req := &lnrpc.GetTransactionsRequest{
			StartHeight: 0,
			EndHeight:   -1,
		}
		txDetails, err := node.GetTransactions(ctxb, req)
		require.NoError(t.t, err)

		for _, tx := range txDetails.Transactions {
			require.Greater(t.t, len(tx.PreviousOutpoints), 0)
		}

		// Lastly, shutdown this Carol so we can move on to the next
		// restoration.
		shutdownAndAssert(net, t, node)
	}

	// Create a closure-factory for building closures that can generate and
	// skip a configurable number of addresses, before finally sending coins
	// to a next generated address. The returned closure will apply the same
	// behavior to both default P2WKH and NP2WKH scopes.
	skipAndSend := func(nskip int) func(*lntest.HarnessNode) {
		return func(node *lntest.HarnessNode) {
			t.t.Helper()

			newP2WKHAddrReq := &lnrpc.NewAddressRequest{
				Type: AddrTypeWitnessPubkeyHash,
			}

			newNP2WKHAddrReq := &lnrpc.NewAddressRequest{
				Type: AddrTypeNestedPubkeyHash,
			}

			newP2TRAddrReq := &lnrpc.NewAddressRequest{
				Type: AddrTypeTaprootPubkey,
			}

			// Generate and skip the number of addresses requested.
			ctxt, cancel := context.WithTimeout(
				ctxb, defaultTimeout,
			)
			defer cancel()
			for i := 0; i < nskip; i++ {
				_, err = node.NewAddress(ctxt, newP2WKHAddrReq)
				require.NoError(t.t, err)

				_, err = node.NewAddress(ctxt, newNP2WKHAddrReq)
				require.NoError(t.t, err)

				_, err = node.NewAddress(ctxt, newP2TRAddrReq)
				require.NoError(t.t, err)
			}

			// Send one BTC to the next P2WKH address.
			net.SendCoins(t.t, btcutil.SatoshiPerBitcoin, node)

			// And another to the next NP2WKH address.
			net.SendCoinsNP2WKH(
				t.t, btcutil.SatoshiPerBitcoin, node,
			)

			// Add another whole coin to the P2TR address.
			net.SendCoinsP2TR(
				t.t, btcutil.SatoshiPerBitcoin, node,
			)
		}
	}

	// Restore Carol with a recovery window of 0. Since no coins have been
	// sent, her balance should be zero.
	//
	// After, one BTC is sent to both her first external P2WKH and NP2WKH
	// addresses.
	restoreCheckBalance(0, 0, 0, skipAndSend(0))

	// Check that restoring without a look-ahead results in having no funds
	// in the wallet, even though they exist on-chain.
	restoreCheckBalance(0, 0, 0, nil)

	// Now, check that using a look-ahead of 1 recovers the balance from
	// the two transactions above. We should also now have 2 UTXOs in the
	// wallet at the end of the recovery attempt.
	//
	// After, we will generate and skip 9 P2WKH, NP2WKH and P2TR addresses,
	// and send another BTC to the subsequent 10th address in each
	// derivation path.
	restoreCheckBalance(3*btcutil.SatoshiPerBitcoin, 3, 1, skipAndSend(9))

	// Check that using a recovery window of 9 does not find the two most
	// recent txns.
	restoreCheckBalance(3*btcutil.SatoshiPerBitcoin, 3, 9, nil)

	// Extending our recovery window to 10 should find the most recent
	// transactions, leaving the wallet with 6 BTC total. We should also
	// learn of the two additional UTXOs created above.
	//
	// After, we will skip 19 more addrs, sending to the 20th address past
	// our last found address, and repeat the same checks.
	restoreCheckBalance(6*btcutil.SatoshiPerBitcoin, 6, 10, skipAndSend(19))

	// Check that recovering with a recovery window of 19 fails to find the
	// most recent transactions.
	restoreCheckBalance(6*btcutil.SatoshiPerBitcoin, 6, 19, nil)

	// Ensure that using a recovery window of 20 succeeds with all UTXOs
	// found and the final balance reflected.

	// After these checks are done, we'll want to make sure we can also
	// recover change address outputs.  This is mainly motivated by a now
	// fixed bug in the wallet in which change addresses could at times be
	// created outside of the default key scopes. Recovery only used to be
	// performed on the default key scopes, so ideally this test case
	// would've caught the bug earlier. Carol has received 9 BTC so far from
	// the miner, we'll send 8 back to ensure all of her UTXOs get spent to
	// avoid fee discrepancies and a change output is formed.
	const minerAmt = 8 * btcutil.SatoshiPerBitcoin
	const finalBalance = 9 * btcutil.SatoshiPerBitcoin
	promptChangeAddr := func(node *lntest.HarnessNode) {
		t.t.Helper()

		minerAddr, err := net.Miner.NewAddress()
		require.NoError(t.t, err)
		ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
		resp, err := node.SendCoins(ctxt, &lnrpc.SendCoinsRequest{
			Addr:   minerAddr.String(),
			Amount: minerAmt,
		})
		require.NoError(t.t, err)
		txid, err := waitForTxInMempool(
			net.Miner.Client, minerMempoolTimeout,
		)
		require.NoError(t.t, err)
		require.Equal(t.t, txid.String(), resp.Txid)

		block := mineBlocks(t, net, 1, 1)[0]
		assertTxInBlock(t, block, txid)
	}
	restoreCheckBalance(finalBalance, 9, 20, promptChangeAddr)

	// We should expect a static fee of 50100 satoshis for spending 9 inputs
	// (3 P2WPKH, 3 NP2WPKH, 3 P2TR) to two P2WPKH outputs. Carol should
	// therefore only have one UTXO present (the change output) of
	// 9 - 8 - fee BTC.
	const fee = 50100
	restoreCheckBalance(finalBalance-minerAmt-fee, 1, 21, nil)

	// Last of all, make sure we can also restore a node from the extended
	// master root key directly instead of the seed.
	var seedMnemonic aezeed.Mnemonic
	copy(seedMnemonic[:], mnemonic)
	cipherSeed, err := seedMnemonic.ToCipherSeed(password)
	require.NoError(t.t, err)
	extendedRootKey, err := hdkeychain.NewMaster(
		cipherSeed.Entropy[:], harnessNetParams,
	)
	require.NoError(t.t, err)
	rootKey = extendedRootKey.String()
	mnemonic = nil

	restoreCheckBalance(finalBalance-minerAmt-fee, 1, 21, nil)
}
<|endoftext|>


<commit_before>package chess

import (
	"bytes"
	"encoding/binary"
	"errors"
	"strconv"
	"strings"
)

// A Board represents a chess board and its relationship between squares and pieces.
type Board struct {
	bbWhiteKing   bitboard
	bbWhiteQueen  bitboard
	bbWhiteRook   bitboard
	bbWhiteBishop bitboard
	bbWhiteKnight bitboard
	bbWhitePawn   bitboard
	bbBlackKing   bitboard
	bbBlackQueen  bitboard
	bbBlackRook   bitboard
	bbBlackBishop bitboard
	bbBlackKnight bitboard
	bbBlackPawn   bitboard
	whiteSqs      bitboard
	blackSqs      bitboard
	emptySqs      bitboard
	whiteKingSq   Square
	blackKingSq   Square
}

// SquareMap returns a mapping of squares to pieces.  A square is only added to the map if it is occupied.
func (b *Board) SquareMap() map[Square]Piece {
	m := map[Square]Piece{}
	for sq := 0; sq < numOfSquaresInBoard; sq++ {
		p := b.Piece(Square(sq))
		if p != NoPiece {
			m[Square(sq)] = p
		}
	}
	return m
}

// Draw returns visual representation of the board useful for debugging.
func (b *Board) Draw() string {
	s := "\n A B C D E F G H\n"
	for r := 7; r >= 0; r-- {
		s += Rank(r).String()
		for f := 0; f < numOfSquaresInRow; f++ {
			p := b.Piece(getSquare(File(f), Rank(r)))
			if p == NoPiece {
				s += "-"
			} else {
				s += p.String()
			}
			s += " "
		}
		s += "\n"
	}
	return s
}

// String implements the fmt.Stringer interface and returns
// a string in the FEN board format: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR
func (b *Board) String() string {
	fen := ""
	for r := 7; r >= 0; r-- {
		for f := 0; f < numOfSquaresInRow; f++ {
			sq := getSquare(File(f), Rank(r))
			p := b.Piece(sq)
			if p != NoPiece {
				fen += p.getFENChar()
			} else {
				fen += "1"
			}
		}
		if r != 0 {
			fen += "/"
		}
	}
	for i := 8; i > 1; i-- {
		repeatStr := strings.Repeat("1", i)
		countStr := strconv.Itoa(i)
		fen = strings.Replace(fen, repeatStr, countStr, -1)
	}
	return fen
}

// Piece returns the piece for the given square.
func (b *Board) Piece(sq Square) Piece {
	for _, p := range allPieces {
		bb := b.bbForPiece(p)
		if bb.Occupied(sq) {
			return p
		}
	}
	return NoPiece
}

// MarshalText implements the encoding.TextMarshaler interface and returns
// a string in the FEN board format: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR
func (b *Board) MarshalText() (text []byte, err error) {
	return []byte(b.String()), nil
}

// UnmarshalText implements the encoding.TextUnarshaler interface and takes
// a string in the FEN board format: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR
func (b *Board) UnmarshalText(text []byte) error {
	cp, err := fenBoard(string(text))
	if err != nil {
		return err
	}
	*b = *cp
	return nil
}

// MarshalBinary implements the encoding.BinaryMarshaler interface and returns
// the bitboard representations as a array of bytes.  Bitboads are encoded
// in the following order: WhiteKing, WhiteQueen, WhiteRook, WhiteBishop, WhiteKnight
// WhitePawn, BlackKing, BlackQueen, BlackRook, BlackBishop, BlackKnight, BlackPawn
func (b *Board) MarshalBinary() (data []byte, err error) {
	bbs := []bitboard{b.bbWhiteKing, b.bbWhiteQueen, b.bbWhiteRook, b.bbWhiteBishop, b.bbWhiteKnight, b.bbWhitePawn,
		b.bbBlackKing, b.bbBlackQueen, b.bbBlackRook, b.bbBlackBishop, b.bbBlackKnight, b.bbBlackPawn}
	buf := new(bytes.Buffer)
	err = binary.Write(buf, binary.BigEndian, bbs)
	return buf.Bytes(), err
}

// UnmarshalBinary implements the encoding.BinaryUnmarshaler interface and parses
// the bitboard representations as a array of bytes.  Bitboads are decoded
// in the following order: WhiteKing, WhiteQueen, WhiteRook, WhiteBishop, WhiteKnight
// WhitePawn, BlackKing, BlackQueen, BlackRook, BlackBishop, BlackKnight, BlackPawn
func (b *Board) UnmarshalBinary(data []byte) error {
	if len(data) != 96 {
		return errors.New("chess: invalid number of bytes for board unmarshal binary")
	}
	b.bbWhiteKing = bitboard(binary.BigEndian.Uint64(data[:8]))
	b.bbWhiteQueen = bitboard(binary.BigEndian.Uint64(data[8:16]))
	b.bbWhiteRook = bitboard(binary.BigEndian.Uint64(data[16:24]))
	b.bbWhiteBishop = bitboard(binary.BigEndian.Uint64(data[24:32]))
	b.bbWhiteKnight = bitboard(binary.BigEndian.Uint64(data[32:40]))
	b.bbWhitePawn = bitboard(binary.BigEndian.Uint64(data[40:48]))
	b.bbBlackKing = bitboard(binary.BigEndian.Uint64(data[48:56]))
	b.bbBlackQueen = bitboard(binary.BigEndian.Uint64(data[56:64]))
	b.bbBlackRook = bitboard(binary.BigEndian.Uint64(data[64:72]))
	b.bbBlackBishop = bitboard(binary.BigEndian.Uint64(data[72:80]))
	b.bbBlackKnight = bitboard(binary.BigEndian.Uint64(data[80:88]))
	b.bbBlackPawn = bitboard(binary.BigEndian.Uint64(data[88:96]))
	b.calcConvienceBBs(nil)
	return nil
}

func newBoard(m map[Square]Piece) *Board {
	b := &Board{}
	for _, p1 := range allPieces {
		bm := map[Square]bool{}
		for sq, p2 := range m {
			if p1 == p2 {
				bm[sq] = true
			}
		}
		bb := newBitboard(bm)
		b.setBBForPiece(p1, bb)
	}
	b.calcConvienceBBs(nil)
	return b
}

func (b *Board) update(m *Move) {
	p1 := b.Piece(m.s1)
	s1BB := bbForSquare(m.s1)
	s2BB := bbForSquare(m.s2)

	// move s1 piece to s2
	for _, p := range allPieces {
		bb := b.bbForPiece(p)
		// remove what was at s2
		b.setBBForPiece(p, bb & ^s2BB)
		// move what was at s1 to s2
		if bb.Occupied(m.s1) {
			bb = b.bbForPiece(p)
			b.setBBForPiece(p, (bb & ^s1BB)|s2BB)
		}
	}
	// check promotion
	if m.promo != NoPieceType {
		newPiece := getPiece(m.promo, p1.Color())
		// remove pawn
		bbPawn := b.bbForPiece(p1)
		b.setBBForPiece(p1, bbPawn & ^s2BB)
		// add promo piece
		bbPromo := b.bbForPiece(newPiece)
		b.setBBForPiece(newPiece, bbPromo|s2BB)
	}
	// remove captured en passant piece
	if m.HasTag(EnPassant) {
		if p1.Color() == White {
			b.bbBlackPawn = ^(bbForSquare(m.s2) << 8) & b.bbBlackPawn
		} else {
			b.bbWhitePawn = ^(bbForSquare(m.s2) >> 8) & b.bbWhitePawn
		}
	}
	// move rook for castle
	if p1.Color() == White && m.HasTag(KingSideCastle) {
		b.bbWhiteRook = (b.bbWhiteRook & ^bbForSquare(H1) | bbForSquare(F1))
	} else if p1.Color() == White && m.HasTag(QueenSideCastle) {
		b.bbWhiteRook = (b.bbWhiteRook & ^bbForSquare(A1)) | bbForSquare(D1)
	} else if p1.Color() == Black && m.HasTag(KingSideCastle) {
		b.bbBlackRook = (b.bbBlackRook & ^bbForSquare(H8) | bbForSquare(F8))
	} else if p1.Color() == Black && m.HasTag(QueenSideCastle) {
		b.bbBlackRook = (b.bbBlackRook & ^bbForSquare(A8)) | bbForSquare(D8)
	}
	b.calcConvienceBBs(m)
}

func (b *Board) calcConvienceBBs(m *Move) {
	whiteSqs := b.bbWhiteKing | b.bbWhiteQueen | b.bbWhiteRook | b.bbWhiteBishop | b.bbWhiteKnight | b.bbWhitePawn
	blackSqs := b.bbBlackKing | b.bbBlackQueen | b.bbBlackRook | b.bbBlackBishop | b.bbBlackKnight | b.bbBlackPawn
	emptySqs := ^(whiteSqs | blackSqs)
	b.whiteSqs = whiteSqs
	b.blackSqs = blackSqs
	b.emptySqs = emptySqs
	if m == nil {
		b.whiteKingSq = NoSquare
		b.blackKingSq = NoSquare

		for sq := 0; sq < numOfSquaresInBoard; sq++ {
			sqr := Square(sq)
			if b.bbWhiteKing.Occupied(sqr) {
				b.whiteKingSq = sqr
			} else if b.bbBlackKing.Occupied(sqr) {
				b.blackKingSq = sqr
			}
		}
	} else if m.s1 == b.whiteKingSq {
		b.whiteKingSq = m.s2
	} else if m.s1 == b.blackKingSq {
		b.blackKingSq = m.s2
	}
}

func (b *Board) copy() *Board {
	return &Board{
		whiteSqs:      b.whiteSqs,
		blackSqs:      b.blackSqs,
		emptySqs:      b.emptySqs,
		whiteKingSq:   b.whiteKingSq,
		blackKingSq:   b.blackKingSq,
		bbWhiteKing:   b.bbWhiteKing,
		bbWhiteQueen:  b.bbWhiteQueen,
		bbWhiteRook:   b.bbWhiteRook,
		bbWhiteBishop: b.bbWhiteBishop,
		bbWhiteKnight: b.bbWhiteKnight,
		bbWhitePawn:   b.bbWhitePawn,
		bbBlackKing:   b.bbBlackKing,
		bbBlackQueen:  b.bbBlackQueen,
		bbBlackRook:   b.bbBlackRook,
		bbBlackBishop: b.bbBlackBishop,
		bbBlackKnight: b.bbBlackKnight,
		bbBlackPawn:   b.bbBlackPawn,
	}
}

func (b *Board) isOccupied(sq Square) bool {
	return !b.emptySqs.Occupied(sq)
}

func (b *Board) hasSufficientMaterial() bool {
	// queen, rook, or pawn exist
	if (b.bbWhiteQueen | b.bbWhiteRook | b.bbWhitePawn |
		b.bbBlackQueen | b.bbBlackRook | b.bbBlackPawn) > 0 {
		return true
	}
	// if king is missing then it is a test
	if b.bbWhiteKing == 0 || b.bbBlackKing == 0 {
		return true
	}
	count := map[PieceType]int{}
	pieceMap := b.SquareMap()
	for _, p := range pieceMap {
		count[p.Type()]++
	}
	// 	king versus king
	if count[Bishop] == 0 && count[Knight] == 0 {
		return false
	}
	// king and bishop versus king
	if count[Bishop] == 1 && count[Knight] == 0 {
		return false
	}
	// king and knight versus king
	if count[Bishop] == 0 && count[Knight] == 1 {
		return false
	}
	// king and bishop(s) versus king and bishop(s) with the bishops on the same colour.
	if count[Knight] == 0 {
		whiteCount := 0
		blackCount := 0
		for sq, p := range pieceMap {
			if p.Type() == Bishop {
				switch sq.color() {
				case White:
					whiteCount++
				case Black:
					blackCount++
				}
			}
		}
		if whiteCount == 0 || blackCount == 0 {
			return false
		}
	}
	return true
}

func (b *Board) bbForPiece(p Piece) bitboard {
	switch p {
	case WhiteKing:
		return b.bbWhiteKing
	case WhiteQueen:
		return b.bbWhiteQueen
	case WhiteRook:
		return b.bbWhiteRook
	case WhiteBishop:
		return b.bbWhiteBishop
	case WhiteKnight:
		return b.bbWhiteKnight
	case WhitePawn:
		return b.bbWhitePawn
	case BlackKing:
		return b.bbBlackKing
	case BlackQueen:
		return b.bbBlackQueen
	case BlackRook:
		return b.bbBlackRook
	case BlackBishop:
		return b.bbBlackBishop
	case BlackKnight:
		return b.bbBlackKnight
	case BlackPawn:
		return b.bbBlackPawn
	}
	return bitboard(0)
}

func (b *Board) setBBForPiece(p Piece, bb bitboard) {
	switch p {
	case WhiteKing:
		b.bbWhiteKing = bb
	case WhiteQueen:
		b.bbWhiteQueen = bb
	case WhiteRook:
		b.bbWhiteRook = bb
	case WhiteBishop:
		b.bbWhiteBishop = bb
	case WhiteKnight:
		b.bbWhiteKnight = bb
	case WhitePawn:
		b.bbWhitePawn = bb
	case BlackKing:
		b.bbBlackKing = bb
	case BlackQueen:
		b.bbBlackQueen = bb
	case BlackRook:
		b.bbBlackRook = bb
	case BlackBishop:
		b.bbBlackBishop = bb
	case BlackKnight:
		b.bbBlackKnight = bb
	case BlackPawn:
		b.bbBlackPawn = bb
	default:
		panic("HERE")
	}
}
<commit_msg>Return userful panic message (#38)<commit_after>package chess

import (
	"bytes"
	"encoding/binary"
	"errors"
	"strconv"
	"strings"
)

// A Board represents a chess board and its relationship between squares and pieces.
type Board struct {
	bbWhiteKing   bitboard
	bbWhiteQueen  bitboard
	bbWhiteRook   bitboard
	bbWhiteBishop bitboard
	bbWhiteKnight bitboard
	bbWhitePawn   bitboard
	bbBlackKing   bitboard
	bbBlackQueen  bitboard
	bbBlackRook   bitboard
	bbBlackBishop bitboard
	bbBlackKnight bitboard
	bbBlackPawn   bitboard
	whiteSqs      bitboard
	blackSqs      bitboard
	emptySqs      bitboard
	whiteKingSq   Square
	blackKingSq   Square
}

// SquareMap returns a mapping of squares to pieces.  A square is only added to the map if it is occupied.
func (b *Board) SquareMap() map[Square]Piece {
	m := map[Square]Piece{}
	for sq := 0; sq < numOfSquaresInBoard; sq++ {
		p := b.Piece(Square(sq))
		if p != NoPiece {
			m[Square(sq)] = p
		}
	}
	return m
}

// Draw returns visual representation of the board useful for debugging.
func (b *Board) Draw() string {
	s := "\n A B C D E F G H\n"
	for r := 7; r >= 0; r-- {
		s += Rank(r).String()
		for f := 0; f < numOfSquaresInRow; f++ {
			p := b.Piece(getSquare(File(f), Rank(r)))
			if p == NoPiece {
				s += "-"
			} else {
				s += p.String()
			}
			s += " "
		}
		s += "\n"
	}
	return s
}

// String implements the fmt.Stringer interface and returns
// a string in the FEN board format: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR
func (b *Board) String() string {
	fen := ""
	for r := 7; r >= 0; r-- {
		for f := 0; f < numOfSquaresInRow; f++ {
			sq := getSquare(File(f), Rank(r))
			p := b.Piece(sq)
			if p != NoPiece {
				fen += p.getFENChar()
			} else {
				fen += "1"
			}
		}
		if r != 0 {
			fen += "/"
		}
	}
	for i := 8; i > 1; i-- {
		repeatStr := strings.Repeat("1", i)
		countStr := strconv.Itoa(i)
		fen = strings.Replace(fen, repeatStr, countStr, -1)
	}
	return fen
}

// Piece returns the piece for the given square.
func (b *Board) Piece(sq Square) Piece {
	for _, p := range allPieces {
		bb := b.bbForPiece(p)
		if bb.Occupied(sq) {
			return p
		}
	}
	return NoPiece
}

// MarshalText implements the encoding.TextMarshaler interface and returns
// a string in the FEN board format: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR
func (b *Board) MarshalText() (text []byte, err error) {
	return []byte(b.String()), nil
}

// UnmarshalText implements the encoding.TextUnarshaler interface and takes
// a string in the FEN board format: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR
func (b *Board) UnmarshalText(text []byte) error {
	cp, err := fenBoard(string(text))
	if err != nil {
		return err
	}
	*b = *cp
	return nil
}

// MarshalBinary implements the encoding.BinaryMarshaler interface and returns
// the bitboard representations as a array of bytes.  Bitboads are encoded
// in the following order: WhiteKing, WhiteQueen, WhiteRook, WhiteBishop, WhiteKnight
// WhitePawn, BlackKing, BlackQueen, BlackRook, BlackBishop, BlackKnight, BlackPawn
func (b *Board) MarshalBinary() (data []byte, err error) {
	bbs := []bitboard{b.bbWhiteKing, b.bbWhiteQueen, b.bbWhiteRook, b.bbWhiteBishop, b.bbWhiteKnight, b.bbWhitePawn,
		b.bbBlackKing, b.bbBlackQueen, b.bbBlackRook, b.bbBlackBishop, b.bbBlackKnight, b.bbBlackPawn}
	buf := new(bytes.Buffer)
	err = binary.Write(buf, binary.BigEndian, bbs)
	return buf.Bytes(), err
}

// UnmarshalBinary implements the encoding.BinaryUnmarshaler interface and parses
// the bitboard representations as a array of bytes.  Bitboads are decoded
// in the following order: WhiteKing, WhiteQueen, WhiteRook, WhiteBishop, WhiteKnight
// WhitePawn, BlackKing, BlackQueen, BlackRook, BlackBishop, BlackKnight, BlackPawn
func (b *Board) UnmarshalBinary(data []byte) error {
	if len(data) != 96 {
		return errors.New("chess: invalid number of bytes for board unmarshal binary")
	}
	b.bbWhiteKing = bitboard(binary.BigEndian.Uint64(data[:8]))
	b.bbWhiteQueen = bitboard(binary.BigEndian.Uint64(data[8:16]))
	b.bbWhiteRook = bitboard(binary.BigEndian.Uint64(data[16:24]))
	b.bbWhiteBishop = bitboard(binary.BigEndian.Uint64(data[24:32]))
	b.bbWhiteKnight = bitboard(binary.BigEndian.Uint64(data[32:40]))
	b.bbWhitePawn = bitboard(binary.BigEndian.Uint64(data[40:48]))
	b.bbBlackKing = bitboard(binary.BigEndian.Uint64(data[48:56]))
	b.bbBlackQueen = bitboard(binary.BigEndian.Uint64(data[56:64]))
	b.bbBlackRook = bitboard(binary.BigEndian.Uint64(data[64:72]))
	b.bbBlackBishop = bitboard(binary.BigEndian.Uint64(data[72:80]))
	b.bbBlackKnight = bitboard(binary.BigEndian.Uint64(data[80:88]))
	b.bbBlackPawn = bitboard(binary.BigEndian.Uint64(data[88:96]))
	b.calcConvienceBBs(nil)
	return nil
}

func newBoard(m map[Square]Piece) *Board {
	b := &Board{}
	for _, p1 := range allPieces {
		bm := map[Square]bool{}
		for sq, p2 := range m {
			if p1 == p2 {
				bm[sq] = true
			}
		}
		bb := newBitboard(bm)
		b.setBBForPiece(p1, bb)
	}
	b.calcConvienceBBs(nil)
	return b
}

func (b *Board) update(m *Move) {
	p1 := b.Piece(m.s1)
	s1BB := bbForSquare(m.s1)
	s2BB := bbForSquare(m.s2)

	// move s1 piece to s2
	for _, p := range allPieces {
		bb := b.bbForPiece(p)
		// remove what was at s2
		b.setBBForPiece(p, bb & ^s2BB)
		// move what was at s1 to s2
		if bb.Occupied(m.s1) {
			bb = b.bbForPiece(p)
			b.setBBForPiece(p, (bb & ^s1BB)|s2BB)
		}
	}
	// check promotion
	if m.promo != NoPieceType {
		newPiece := getPiece(m.promo, p1.Color())
		// remove pawn
		bbPawn := b.bbForPiece(p1)
		b.setBBForPiece(p1, bbPawn & ^s2BB)
		// add promo piece
		bbPromo := b.bbForPiece(newPiece)
		b.setBBForPiece(newPiece, bbPromo|s2BB)
	}
	// remove captured en passant piece
	if m.HasTag(EnPassant) {
		if p1.Color() == White {
			b.bbBlackPawn = ^(bbForSquare(m.s2) << 8) & b.bbBlackPawn
		} else {
			b.bbWhitePawn = ^(bbForSquare(m.s2) >> 8) & b.bbWhitePawn
		}
	}
	// move rook for castle
	if p1.Color() == White && m.HasTag(KingSideCastle) {
		b.bbWhiteRook = (b.bbWhiteRook & ^bbForSquare(H1) | bbForSquare(F1))
	} else if p1.Color() == White && m.HasTag(QueenSideCastle) {
		b.bbWhiteRook = (b.bbWhiteRook & ^bbForSquare(A1)) | bbForSquare(D1)
	} else if p1.Color() == Black && m.HasTag(KingSideCastle) {
		b.bbBlackRook = (b.bbBlackRook & ^bbForSquare(H8) | bbForSquare(F8))
	} else if p1.Color() == Black && m.HasTag(QueenSideCastle) {
		b.bbBlackRook = (b.bbBlackRook & ^bbForSquare(A8)) | bbForSquare(D8)
	}
	b.calcConvienceBBs(m)
}

func (b *Board) calcConvienceBBs(m *Move) {
	whiteSqs := b.bbWhiteKing | b.bbWhiteQueen | b.bbWhiteRook | b.bbWhiteBishop | b.bbWhiteKnight | b.bbWhitePawn
	blackSqs := b.bbBlackKing | b.bbBlackQueen | b.bbBlackRook | b.bbBlackBishop | b.bbBlackKnight | b.bbBlackPawn
	emptySqs := ^(whiteSqs | blackSqs)
	b.whiteSqs = whiteSqs
	b.blackSqs = blackSqs
	b.emptySqs = emptySqs
	if m == nil {
		b.whiteKingSq = NoSquare
		b.blackKingSq = NoSquare

		for sq := 0; sq < numOfSquaresInBoard; sq++ {
			sqr := Square(sq)
			if b.bbWhiteKing.Occupied(sqr) {
				b.whiteKingSq = sqr
			} else if b.bbBlackKing.Occupied(sqr) {
				b.blackKingSq = sqr
			}
		}
	} else if m.s1 == b.whiteKingSq {
		b.whiteKingSq = m.s2
	} else if m.s1 == b.blackKingSq {
		b.blackKingSq = m.s2
	}
}

func (b *Board) copy() *Board {
	return &Board{
		whiteSqs:      b.whiteSqs,
		blackSqs:      b.blackSqs,
		emptySqs:      b.emptySqs,
		whiteKingSq:   b.whiteKingSq,
		blackKingSq:   b.blackKingSq,
		bbWhiteKing:   b.bbWhiteKing,
		bbWhiteQueen:  b.bbWhiteQueen,
		bbWhiteRook:   b.bbWhiteRook,
		bbWhiteBishop: b.bbWhiteBishop,
		bbWhiteKnight: b.bbWhiteKnight,
		bbWhitePawn:   b.bbWhitePawn,
		bbBlackKing:   b.bbBlackKing,
		bbBlackQueen:  b.bbBlackQueen,
		bbBlackRook:   b.bbBlackRook,
		bbBlackBishop: b.bbBlackBishop,
		bbBlackKnight: b.bbBlackKnight,
		bbBlackPawn:   b.bbBlackPawn,
	}
}

func (b *Board) isOccupied(sq Square) bool {
	return !b.emptySqs.Occupied(sq)
}

func (b *Board) hasSufficientMaterial() bool {
	// queen, rook, or pawn exist
	if (b.bbWhiteQueen | b.bbWhiteRook | b.bbWhitePawn |
		b.bbBlackQueen | b.bbBlackRook | b.bbBlackPawn) > 0 {
		return true
	}
	// if king is missing then it is a test
	if b.bbWhiteKing == 0 || b.bbBlackKing == 0 {
		return true
	}
	count := map[PieceType]int{}
	pieceMap := b.SquareMap()
	for _, p := range pieceMap {
		count[p.Type()]++
	}
	// 	king versus king
	if count[Bishop] == 0 && count[Knight] == 0 {
		return false
	}
	// king and bishop versus king
	if count[Bishop] == 1 && count[Knight] == 0 {
		return false
	}
	// king and knight versus king
	if count[Bishop] == 0 && count[Knight] == 1 {
		return false
	}
	// king and bishop(s) versus king and bishop(s) with the bishops on the same colour.
	if count[Knight] == 0 {
		whiteCount := 0
		blackCount := 0
		for sq, p := range pieceMap {
			if p.Type() == Bishop {
				switch sq.color() {
				case White:
					whiteCount++
				case Black:
					blackCount++
				}
			}
		}
		if whiteCount == 0 || blackCount == 0 {
			return false
		}
	}
	return true
}

func (b *Board) bbForPiece(p Piece) bitboard {
	switch p {
	case WhiteKing:
		return b.bbWhiteKing
	case WhiteQueen:
		return b.bbWhiteQueen
	case WhiteRook:
		return b.bbWhiteRook
	case WhiteBishop:
		return b.bbWhiteBishop
	case WhiteKnight:
		return b.bbWhiteKnight
	case WhitePawn:
		return b.bbWhitePawn
	case BlackKing:
		return b.bbBlackKing
	case BlackQueen:
		return b.bbBlackQueen
	case BlackRook:
		return b.bbBlackRook
	case BlackBishop:
		return b.bbBlackBishop
	case BlackKnight:
		return b.bbBlackKnight
	case BlackPawn:
		return b.bbBlackPawn
	}
	return bitboard(0)
}

func (b *Board) setBBForPiece(p Piece, bb bitboard) {
	switch p {
	case WhiteKing:
		b.bbWhiteKing = bb
	case WhiteQueen:
		b.bbWhiteQueen = bb
	case WhiteRook:
		b.bbWhiteRook = bb
	case WhiteBishop:
		b.bbWhiteBishop = bb
	case WhiteKnight:
		b.bbWhiteKnight = bb
	case WhitePawn:
		b.bbWhitePawn = bb
	case BlackKing:
		b.bbBlackKing = bb
	case BlackQueen:
		b.bbBlackQueen = bb
	case BlackRook:
		b.bbBlackRook = bb
	case BlackBishop:
		b.bbBlackBishop = bb
	case BlackKnight:
		b.bbBlackKnight = bb
	case BlackPawn:
		b.bbBlackPawn = bb
	default:
		panic("invalid piece")
	}
}
<|endoftext|>


<commit_before>package teamcity

import (
	"strconv"
	"strings"
	"time"
)

const (
	dateFormat = "20060102T150405-0700"
)

// Builds is a list of builds
type Builds struct {
	Count    int     `json:"count,omitempty"`
	NextHref string  `json:"nextHref,omitempty"`
	Builds   []Build `json:"build"`
}

// Build is an instance of a stage in the build chain for a given project
type Build struct {
	Id              int             `json:"id,omitempty"`
	Number          string          `json:"number,omitempty"`
	BuildTypeId     string          `json:"buildTypeId,omitempty"`
	BuildType       BuildType       `json:"buildType,omitempty"`
	Status          string          `json:"status,omitempty"`
	State           string          `json:"state,omitempty"`
	Href            string          `json:"href,omitempty"`
	StatusText      string          `json:"statusText,omitempty"`
	QueuedDate      Time            `json:"queuedDate,omitempty"`
	StartDate       Time            `json:"startDate,omitempty"`
	FinishDate      Time            `json:"finishDate,omitempty"`
	Changes         Changes         `json:"changes,omitempty"`
	LastChanges     Changes         `json:"lastChanges,omitempty"`
	Triggered       Triggered       `json:"triggered,omitempty"`
	Comment         Comment         `json:"comment,omitempty"`
	Properties      Params          `json:"properties,omitempty"`
	WebUrl          string          `json:"webUrl,omitempty"`
	BuildStatistics BuildStatistics `json:"statistics,omitempty"`
}

// BuildType is a type of Build
type BuildType struct {
	Id                   string                `json:"id,omitempty"`
	Name                 string                `json:"name,omitempty"`
	SnapshotDependencies *SnapshotDependencies `json:"snapshot-dependencies,omitempty"`
	Project              *Project              `json:"project,omitempty"`
	VcsRootEntries       *VcsRootEntries       `json:"vcs-root-entries"`
	Template             *BuildType            `json:"template,omitempty"`
	Parameters           Params                `json:"parameters,omitempty"`
}

// BuildTypes is a container for a list of BuildType's
type BuildTypes struct {
	BuildTypes []BuildType `json:"buildType,omitempty"`
}

// Dependency is a build type's artifact or snapshot dependency
type Dependency struct {
	Id              string        `json:"id,omitempty"`
	Type            string        `json:"type,omitempty"`
	SourceBuildType BuildType     `json:"source-buildType,omitempty"`
	PropertyList    *PropertyList `json:"properties,omitempty"`
}

type ArtifactDependencies struct {
	ArtifactDependencies []Dependency `json:"artifact-dependency"`
}

// PropertyList is a list of name-value attributes describing some entity.
type PropertyList struct {
	Count      int        `json:"count"`
	Properties []Property `json:"property"`
}

func NewPropertyList(m map[string]string) *PropertyList {
	var props []Property
	for k, v := range m {
		props = append(props, Property{Name: k, Value: v})
	}
	return &PropertyList{Count: len(props), Properties: props}
}

// Value returns the named property's value, or empty string if not found.
func (pl *PropertyList) Value(name string) string {
	if pl == nil {
		return ""
	}
	for _, v := range pl.Properties {
		if v.Name == name {
			return v.Value
		}
	}
	return ""
}

// Bool returns the named property's boolean value, or false if not found.
func (pl *PropertyList) Bool(name string) bool {
	if pl == nil {
		return false
	}
	var val = pl.Value(name)
	b, err := strconv.ParseBool(val)
	if err != nil {
		return false
	}
	return b
}

// Comment is a description for a Build instance
type Comment struct {
	Text string `json:"text"`
}

// Changes are the list of changes that corresponds to a certain build
type Changes struct {
	Changes []Change `json:"change"`
}

// GetChange returns the most relevant Change describing the build, prioritizing
// Build.Changes over Build.LastChanges out of preference for changes to non-TeamCity repos
func (b *Build) GetChange() Change {
	if len(b.Changes.Changes) > 0 {
		return b.Changes.Changes[0]
	}
	if len(b.LastChanges.Changes) > 0 {
		return b.LastChanges.Changes[0]
	}
	return Change{}
}

// Change is an individual change in a group that corresponds to a certain build
type Change struct {
	Id       int    `json:"id,omitempty"`
	Version  string `json:"version,omitempty"`
	Username string `json:"username,omitempty"`
	Date     Time   `json:"date,omitempty"`
	Comment  string `json:"comment,omitempty"`
}

// GetShortVersion returns the first 8 characters of the change version
func (c *Change) GetShortVersion() string {
	var v string
	if len(c.Version) >= 8 {
		v = c.Version[:8]
	}
	return v
}

// BuildsByDate is an interface for sorting a Build array by Date
type BuildsByDate []Build

// Functions for using Golang "sort" package
func (c BuildsByDate) Len() int      { return len(c) }
func (c BuildsByDate) Swap(i, j int) { c[i], c[j] = c[j], c[i] }
func (c BuildsByDate) Less(i, j int) bool {
	timeA := time.Time(c[i].Triggered.Date)
	timeB := time.Time(c[j].Triggered.Date)
	return timeA.Before(timeB)
}

// SnapshotDependencies is a container for SnapshotDependency's
type SnapshotDependencies struct {
	SnapshotDependencies []SnapshotDependency `json:"snapshot-dependency,omitempty"`
}

// SnapshotDependency relates a build type to its source build type
type SnapshotDependency struct {
	SourceBuildType BuildType `json:"source-buildType,omitempty"`
}

// VcsRootEntries is a list of VcsRootEntry
type VcsRootEntries struct {
	VcsRootEntries []VcsRootEntry `json:"vcs-root-entry,omitempty"`
}

// VcsRootEntry is a version control system entry for a build type
type VcsRootEntry struct {
	Id      string  `json:"id,omitempty"`
	VcsRoot VcsRoot `json:"vcs-root,omitempty"`
}

// VcsRoot is a the id, name and properties of a version control system
type VcsRoot struct {
	Id           string        `json:"id,omitempty"`
	Name         string        `json:"name,omitempty"`
	PropertyList *PropertyList `json:"properties,omitempty"`
}

type BuildStatistics struct {
	StatisticsEntries []StatisticsEntry `json:"property,omitempty"`
}

// StatisticsEntry gives a name-value pair of a statistic
type StatisticsEntry struct {
	Name  string `json:"name,omitempty"`
	Value string `json:"value,omitempty"`
}

type Tag struct {
	Name string `json:"name,omitempty"`
}

type Tags struct {
	Tags []Tag `json:"tag,omitempty"`
}

func NewTags(t []string) *Tags {
	tags := []Tag{}
	for _, tag := range t {
		tags = append(tags, Tag{Name: tag})
	}
	return &Tags{Tags: tags}
}

// Triggered describes what triggered a particular build
type Triggered struct {
	Date Time `json:"date,omitempty"`
	User User `json:"user,omitempty"`
}

// User describes a user on TeamCity
type User struct {
	Username string `json:"username,omitempty"`
}

// Time is the date in the format TeamCity provides
type Time time.Time

// UnmarshalJSON unmarshals the time using the TeamCity format
func (t *Time) UnmarshalJSON(buf []byte) error {
	tt, err := time.Parse(dateFormat, strings.Trim(string(buf), `"`))
	if err != nil {
		return err
	}
	*t = Time(tt)
	return nil
}

// MarshalJSON marshals the time using the TeamCity format
func (t Time) MarshalJSON() ([]byte, error) {
	return []byte(`"` + time.Time(t).Format(dateFormat) + `"`), nil
}
<commit_msg>Make Parameters a pointer to a struct in BuildType<commit_after>package teamcity

import (
	"strconv"
	"strings"
	"time"
)

const (
	dateFormat = "20060102T150405-0700"
)

// Builds is a list of builds
type Builds struct {
	Count    int     `json:"count,omitempty"`
	NextHref string  `json:"nextHref,omitempty"`
	Builds   []Build `json:"build"`
}

// Build is an instance of a stage in the build chain for a given project
type Build struct {
	Id              int             `json:"id,omitempty"`
	Number          string          `json:"number,omitempty"`
	BuildTypeId     string          `json:"buildTypeId,omitempty"`
	BuildType       BuildType       `json:"buildType,omitempty"`
	Status          string          `json:"status,omitempty"`
	State           string          `json:"state,omitempty"`
	Href            string          `json:"href,omitempty"`
	StatusText      string          `json:"statusText,omitempty"`
	QueuedDate      Time            `json:"queuedDate,omitempty"`
	StartDate       Time            `json:"startDate,omitempty"`
	FinishDate      Time            `json:"finishDate,omitempty"`
	Changes         Changes         `json:"changes,omitempty"`
	LastChanges     Changes         `json:"lastChanges,omitempty"`
	Triggered       Triggered       `json:"triggered,omitempty"`
	Comment         Comment         `json:"comment,omitempty"`
	Properties      Params          `json:"properties,omitempty"`
	WebUrl          string          `json:"webUrl,omitempty"`
	BuildStatistics BuildStatistics `json:"statistics,omitempty"`
}

// BuildType is a type of Build
type BuildType struct {
	Id                   string                `json:"id,omitempty"`
	Name                 string                `json:"name,omitempty"`
	SnapshotDependencies *SnapshotDependencies `json:"snapshot-dependencies,omitempty"`
	Project              *Project              `json:"project,omitempty"`
	VcsRootEntries       *VcsRootEntries       `json:"vcs-root-entries"`
	Template             *BuildType            `json:"template,omitempty"`
	Parameters           *Params               `json:"parameters,omitempty"`
}

// BuildTypes is a container for a list of BuildType's
type BuildTypes struct {
	BuildTypes []BuildType `json:"buildType,omitempty"`
}

// Dependency is a build type's artifact or snapshot dependency
type Dependency struct {
	Id              string        `json:"id,omitempty"`
	Type            string        `json:"type,omitempty"`
	SourceBuildType BuildType     `json:"source-buildType,omitempty"`
	PropertyList    *PropertyList `json:"properties,omitempty"`
}

type ArtifactDependencies struct {
	ArtifactDependencies []Dependency `json:"artifact-dependency"`
}

// PropertyList is a list of name-value attributes describing some entity.
type PropertyList struct {
	Count      int        `json:"count"`
	Properties []Property `json:"property"`
}

func NewPropertyList(m map[string]string) *PropertyList {
	var props []Property
	for k, v := range m {
		props = append(props, Property{Name: k, Value: v})
	}
	return &PropertyList{Count: len(props), Properties: props}
}

// Value returns the named property's value, or empty string if not found.
func (pl *PropertyList) Value(name string) string {
	if pl == nil {
		return ""
	}
	for _, v := range pl.Properties {
		if v.Name == name {
			return v.Value
		}
	}
	return ""
}

// Bool returns the named property's boolean value, or false if not found.
func (pl *PropertyList) Bool(name string) bool {
	if pl == nil {
		return false
	}
	var val = pl.Value(name)
	b, err := strconv.ParseBool(val)
	if err != nil {
		return false
	}
	return b
}

// Comment is a description for a Build instance
type Comment struct {
	Text string `json:"text"`
}

// Changes are the list of changes that corresponds to a certain build
type Changes struct {
	Changes []Change `json:"change"`
}

// GetChange returns the most relevant Change describing the build, prioritizing
// Build.Changes over Build.LastChanges out of preference for changes to non-TeamCity repos
func (b *Build) GetChange() Change {
	if len(b.Changes.Changes) > 0 {
		return b.Changes.Changes[0]
	}
	if len(b.LastChanges.Changes) > 0 {
		return b.LastChanges.Changes[0]
	}
	return Change{}
}

// Change is an individual change in a group that corresponds to a certain build
type Change struct {
	Id       int    `json:"id,omitempty"`
	Version  string `json:"version,omitempty"`
	Username string `json:"username,omitempty"`
	Date     Time   `json:"date,omitempty"`
	Comment  string `json:"comment,omitempty"`
}

// GetShortVersion returns the first 8 characters of the change version
func (c *Change) GetShortVersion() string {
	var v string
	if len(c.Version) >= 8 {
		v = c.Version[:8]
	}
	return v
}

// BuildsByDate is an interface for sorting a Build array by Date
type BuildsByDate []Build

// Functions for using Golang "sort" package
func (c BuildsByDate) Len() int      { return len(c) }
func (c BuildsByDate) Swap(i, j int) { c[i], c[j] = c[j], c[i] }
func (c BuildsByDate) Less(i, j int) bool {
	timeA := time.Time(c[i].Triggered.Date)
	timeB := time.Time(c[j].Triggered.Date)
	return timeA.Before(timeB)
}

// SnapshotDependencies is a container for SnapshotDependency's
type SnapshotDependencies struct {
	SnapshotDependencies []SnapshotDependency `json:"snapshot-dependency,omitempty"`
}

// SnapshotDependency relates a build type to its source build type
type SnapshotDependency struct {
	SourceBuildType BuildType `json:"source-buildType,omitempty"`
}

// VcsRootEntries is a list of VcsRootEntry
type VcsRootEntries struct {
	VcsRootEntries []VcsRootEntry `json:"vcs-root-entry,omitempty"`
}

// VcsRootEntry is a version control system entry for a build type
type VcsRootEntry struct {
	Id      string  `json:"id,omitempty"`
	VcsRoot VcsRoot `json:"vcs-root,omitempty"`
}

// VcsRoot is a the id, name and properties of a version control system
type VcsRoot struct {
	Id           string        `json:"id,omitempty"`
	Name         string        `json:"name,omitempty"`
	PropertyList *PropertyList `json:"properties,omitempty"`
}

type BuildStatistics struct {
	StatisticsEntries []StatisticsEntry `json:"property,omitempty"`
}

// StatisticsEntry gives a name-value pair of a statistic
type StatisticsEntry struct {
	Name  string `json:"name,omitempty"`
	Value string `json:"value,omitempty"`
}

type Tag struct {
	Name string `json:"name,omitempty"`
}

type Tags struct {
	Tags []Tag `json:"tag,omitempty"`
}

func NewTags(t []string) *Tags {
	tags := []Tag{}
	for _, tag := range t {
		tags = append(tags, Tag{Name: tag})
	}
	return &Tags{Tags: tags}
}

// Triggered describes what triggered a particular build
type Triggered struct {
	Date Time `json:"date,omitempty"`
	User User `json:"user,omitempty"`
}

// User describes a user on TeamCity
type User struct {
	Username string `json:"username,omitempty"`
}

// Time is the date in the format TeamCity provides
type Time time.Time

// UnmarshalJSON unmarshals the time using the TeamCity format
func (t *Time) UnmarshalJSON(buf []byte) error {
	tt, err := time.Parse(dateFormat, strings.Trim(string(buf), `"`))
	if err != nil {
		return err
	}
	*t = Time(tt)
	return nil
}

// MarshalJSON marshals the time using the TeamCity format
func (t Time) MarshalJSON() ([]byte, error) {
	return []byte(`"` + time.Time(t).Format(dateFormat) + `"`), nil
}
<|endoftext|>


<commit_before>package lytics

import (
	"net/url"
	"time"
)

const (
	queryEndpoint         = "query/:id"
	queryListEndpoint     = "query"
	queryTestEndpoint     = "query/_test"
	queryValidateEndpoint = "query/_validate"
)

type (
	// Query represents an LQL Statement structure
	Query struct {
		Id      string                `json:"id"`
		Created time.Time             `json:"created"`
		Updated time.Time             `json:"updated"`
		Alias   string                `json:"alias"`
		Table   string                `json:"table"`
		From    string                `json:"from"`
		Text    string                `json:"text"`
		Fields  map[string]QueryField `json:"fields"`
	}
	// A field in a query
	// - very similar to catalog, query fields create catalog fields
	QueryField struct {
		As         string   `json:"as"`
		IsBy       bool     `json:"is_by"`
		Type       string   `json:"type"`
		ShortDesc  string   `json:"shortdesc"`
		LongDesc   string   `json:"longdesc"`
		Identities []string `json:"identities"`
		Froms      []string `json:"froms"`
	}
)

// GetQueries returns a list of all queries associated with this account
// https://www.getlytics.com/developers/rest-api#query
func (l *Client) GetQueries() ([]Query, error) {
	res := ApiResp{}
	data := []Query{}

	// make the request
	err := l.Get(queryListEndpoint, nil, nil, &res, &data)
	if err != nil {
		return data, err
	}

	return data, nil
}

// GetQueries returns a list of all queries associated with this account
// https://www.getlytics.com/developers/rest-api#query
func (l *Client) GetQuery(alias string) (Query, error) {
	res := ApiResp{}
	data := Query{}

	// make the request
	err := l.Get(parseLyticsURL(queryEndpoint, map[string]string{"id": alias}), nil, nil, &res, &data)
	if err != nil {
		return data, err
	}

	return data, nil
}

// GetQueryTest returns the evaluated entity from given query
// https://www.getlytics.com/developers/rest-api#query
func (l *Client) GetQueryTest(qs url.Values, query string) (Entity, error) {
	res := ApiResp{}
	data := Entity{}

	// make the request
	err := l.Post(queryTestEndpoint, qs, query, &res, &data)

	if err != nil {
		return data, err
	}

	return data, nil
}

// PostQueryValidate returns the query and how it is interpreted
// https://www.getlytics.com/developers/rest-api#query
func (l *Client) PostQueryValidate(query string) (Query, error) {
	res := ApiResp{}
	data := Query{}

	// make the request
	err := l.PostType("text/plain", queryValidateEndpoint, nil, query, &res, &data)

	if err != nil {
		return data, err
	}

	return data, nil
}
<commit_msg>added PostQueryValidateSegment method that pass segments=true in query validate request<commit_after>package lytics

import (
	"net/url"
	"time"
)

const (
	queryEndpoint         = "query/:id"
	queryListEndpoint     = "query"
	queryTestEndpoint     = "query/_test"
	queryValidateEndpoint = "query/_validate"
)

type (
	// Query represents an LQL Statement structure
	Query struct {
		Id      string                `json:"id"`
		Created time.Time             `json:"created"`
		Updated time.Time             `json:"updated"`
		Alias   string                `json:"alias"`
		Table   string                `json:"table"`
		From    string                `json:"from"`
		Text    string                `json:"text"`
		Fields  map[string]QueryField `json:"fields"`
	}
	// A field in a query
	// - very similar to catalog, query fields create catalog fields
	QueryField struct {
		As         string   `json:"as"`
		IsBy       bool     `json:"is_by"`
		Type       string   `json:"type"`
		ShortDesc  string   `json:"shortdesc"`
		LongDesc   string   `json:"longdesc"`
		Identities []string `json:"identities"`
		Froms      []string `json:"froms"`
	}
)

// GetQueries returns a list of all queries associated with this account
// https://www.getlytics.com/developers/rest-api#query
func (l *Client) GetQueries() ([]Query, error) {
	res := ApiResp{}
	data := []Query{}

	// make the request
	err := l.Get(queryListEndpoint, nil, nil, &res, &data)
	if err != nil {
		return data, err
	}

	return data, nil
}

// GetQueries returns a list of all queries associated with this account
// https://www.getlytics.com/developers/rest-api#query
func (l *Client) GetQuery(alias string) (Query, error) {
	res := ApiResp{}
	data := Query{}

	// make the request
	err := l.Get(parseLyticsURL(queryEndpoint, map[string]string{"id": alias}), nil, nil, &res, &data)
	if err != nil {
		return data, err
	}

	return data, nil
}

// GetQueryTest returns the evaluated entity from given query
// https://www.getlytics.com/developers/rest-api#query
func (l *Client) GetQueryTest(qs url.Values, query string) (Entity, error) {
	res := ApiResp{}
	data := Entity{}

	// make the request
	err := l.Post(queryTestEndpoint, qs, query, &res, &data)

	if err != nil {
		return data, err
	}

	return data, nil
}

// PostQueryValidate returns the query and how it is interpreted
// https://www.getlytics.com/developers/rest-api#query
func (l *Client) PostQueryValidate(query string) ([]Query, error) {
	res := ApiResp{}
	data := []Query{}

	// make the request
	err := l.PostType("text/plain", queryValidateEndpoint, nil, query, &res, &data)

	if err != nil {
		return data, err
	}

	return data, nil
}

// PostQueryValidateSegment checks query interpretation & validates query against existing segments/ schemas
func (l *Client) PostQueryValidateSegment(query string) ([]Query, error) {
	res := ApiResp{}
	data := []Query{}
	v := url.Values{}
	v.Set("segments", "true")

	// make the request
	err := l.PostType("text/plain", queryValidateEndpoint, v, query, &res, &data)

	if err != nil {
		return data, err
	}

	return data, nil
}
<|endoftext|>


<commit_before>package hrd

import (
	"github.com/101loops/hrd/internal/types"

	ae "appengine"
)

// Query represents a datastore query.
type Query struct {
	inner *types.Query
	ctx   ae.Context
	kind  *Kind
	opts  *types.Opts
}

// newQuery creates a new Query for the passed kind.
// The kind's options are used as default options.
func newQuery(ctx ae.Context, kind *Kind) (ret *Query) {
	return &Query{
		inner: types.NewQuery(kind.name),
		ctx:   ctx,
		kind:  kind,
		opts:  types.DefaultOpts(),
	}
}

func (qry *Query) clone() *Query {
	ret := *qry
	ret.opts = qry.opts.Clone()
	ret.inner = qry.inner.Clone()
	//	if len(qry.log) > 0 {
	//		ret.log = make([]string, len(qry.log))
	//		copy(ret.log, qry.log)
	//	}
	return &ret
}

// NoGlobalCache prevents reading/writing entities from/to memcache.
func (qry *Query) NoGlobalCache() (ret *Query) {
	ret = qry.clone()
	ret.opts.NoGlobalCache = true
	return
}

// Limit returns a derivative Query that has a limit on the number
// of results returned. A negative value means unlimited.
func (qry *Query) Limit(limit int) (ret *Query) {
	ret = qry.clone()
	if limit > 0 {
		//ret.log("LIMIT %v", limit)
	} else {
		limit = -1
		//ret.log("NO LIMIT")
	}
	ret.inner.Limit = limit
	return
}

// NoLimit returns a derivative Query that has no limit on the number
// of results returned.
func (qry *Query) NoLimit() (ret *Query) {
	return qry.Limit(-1)
}

// Ancestor returns a derivative Query with an ancestor filter.
// The ancestor should not be nil.
func (qry *Query) Ancestor(k *Key) (ret *Query) {
	ret = qry.clone()
	//ret.log("ANCESTOR '%v'", k.String())
	ret.inner.Ancestor = k.inner
	return
}

// Project returns a derivative Query that yields only the passed fields.
// It cannot be used in a keys-only query.
func (qry *Query) Project(fields ...string) (ret *Query) {
	ret = qry.clone()
	//ret.log("PROJECT '%v'", strings.Join(fields, "', '"))
	ret.inner.Projection = append([]string(nil), fields...)
	ret.inner.TypeOf = types.ProjectQuery
	return
}

// EventualConsistency returns a derivative query that returns eventually
// consistent results. It only has an effect on ancestor queries.
func (qry *Query) EventualConsistency() (ret *Query) {
	ret = qry.clone()
	//ret.log("EVENTUAL CONSISTENCY")
	ret.inner.Eventual = true
	return
}

// Start returns a derivative Query with the passed start point.
func (qry *Query) Start(c string) (ret *Query) {
	ret = qry.clone()
	//ret.log("START CURSOR")
	ret.inner.Start = c
	return
}

// End returns a derivative Query with the passed end point.
func (qry *Query) End(c string) (ret *Query) {
	ret = qry.clone()
	//ret.log("END CURSOR")
	ret.inner.End = c
	return
}

// Offset returns a derivative Query that has an offset of how many keys
// to skip over before returning results. A negative value is invalid.
func (qry *Query) Offset(off int) (ret *Query) {
	ret = qry.clone()
	//ret.log("OFFSET %v", off)
	ret.inner.Offset = off
	return
}

// OrderAsc returns a derivative Query with a field-based sort order, ascending.
// Orders are applied in the order they are added.
func (qry *Query) OrderAsc(s string) (ret *Query) {
	ret = qry.clone()
	//ret.log("ORDER ASC %v", s)
	ret.inner.Order = append(ret.inner.Order, types.Order{FieldName: s, Descending: false})
	return
}

// OrderDesc returns a derivative Query with a field-based sort order, descending.
// Orders are applied in the order they are added.
func (qry *Query) OrderDesc(s string) (ret *Query) {
	ret = qry.clone()
	//ret.log("ORDER DESC %v", s)
	ret.inner.Order = append(ret.inner.Order, types.Order{FieldName: s, Descending: true})
	return
}

// Distinct returns a derivative query that yields de-duplicated entities with
// respect to the set of projected fields. It is only used for projection
// queries.
func (qry *Query) Distinct() (ret *Query) {
	ret = qry.clone()
	ret.inner.Distinct = true
	return ret
}

// Filter returns a derivative Query with a field-based filter.
// The filterStr argument must be a field name followed by optional space,
// followed by an operator, one of ">", "<", ">=", "<=", or "=".
// Fields are compared against the provided value using the operator.
// Multiple filters are AND'ed together.
func (qry *Query) Filter(q string, val interface{}) (ret *Query) {
	ret = qry.clone()
	ret.inner.Filter = append(ret.inner.Filter, types.Filter{Filter: q, Value: val})
	//ret.log("FILTER '%v %v'", q, val)
	return
}

// GetCount returns the number of results for the query.
func (qry *Query) GetCount() (int, error) {
	//qry.log("COUNT")
	//qry.ctx.Infof(qry.getLog())

	return dsCount(qry.ctx, qry.inner)
}

// GetKeys executes the query as keys-only: No entities are retrieved, just their keys.
func (qry *Query) GetKeys() ([]*Key, string, error) {
	keysQry := qry.clone()
	//keysQry.log("KEYS-ONLY")
	keysQry.inner.TypeOf = types.KeysOnlyQuery

	it := keysQry.Run()

	keys, err := it.GetAll(nil)
	if err != nil {
		return nil, "", err
	}

	cursor, err := it.Cursor()
	return keys, cursor, err
}

// GetAll runs the query and writes the entities to the passed destination.
//
// Note that, if not manually disabled, queries for more than 1 item use
// a "hybrid query". This means that first a keys-only query is executed
// and then the keys are used to lookup the local and global cache as well
// as the datastore eventually. For a warm cache this usually is
// faster and cheaper than the regular query.
func (qry *Query) GetAll(dsts interface{}) ([]*Key, string, error) {
	useHybridQry := qry.inner.Limit != 1 && qry.inner.TypeOf == types.FullQuery && !qry.opts.NoGlobalCache
	if useHybridQry {
		return qry.getAllByHybrid(dsts)
	}

	it := qry.Run()
	keys, err := refactorit.GetAll(dsts)
	if err != nil {
		return nil, "", err
	}
	cursor, err := it.Cursor()
	return keys, cursor, err
}

func (qry *Query) getAllByHybrid(dsts interface{}) ([]*Key, string, error) {
	keys, cursor, err := qry.GetKeys()
	if err == nil && len(keys) > 0 {
		keys, err = newLoader(qry.ctx, qry.kind).Keys(keys).GetAll(dsts)
	}
	return keys, cursor, err
}

// GetFirst executes the query and writes the result's first entity
// to the passed destination.
func (qry *Query) GetFirst(dst interface{}) (*Key, error) {
	return qry.Run().GetOne(dst)
}

// Run executes the query and returns an Iterator.
func (qry *Query) Run() *Iterator {
	//qry.ctx.Infof(qry.getLog())
	return newIterator(qry)
}

//func (qry *Query) log(s string, values ...interface{}) {
//	qry.log = append(qry.log, fmt.Sprintf(s, values...))
//}
//
//func (qry *Query) getLog() string {
//	return fmt.Sprintf("running query \"%v\"", strings.Join(qry.log, " | "))
//}
<commit_msg>fix typo<commit_after>package hrd

import (
	"github.com/101loops/hrd/internal/types"

	ae "appengine"
)

// Query represents a datastore query.
type Query struct {
	inner *types.Query
	ctx   ae.Context
	kind  *Kind
	opts  *types.Opts
}

// newQuery creates a new Query for the passed kind.
// The kind's options are used as default options.
func newQuery(ctx ae.Context, kind *Kind) (ret *Query) {
	return &Query{
		inner: types.NewQuery(kind.name),
		ctx:   ctx,
		kind:  kind,
		opts:  types.DefaultOpts(),
	}
}

func (qry *Query) clone() *Query {
	ret := *qry
	ret.opts = qry.opts.Clone()
	ret.inner = qry.inner.Clone()
	//	if len(qry.log) > 0 {
	//		ret.log = make([]string, len(qry.log))
	//		copy(ret.log, qry.log)
	//	}
	return &ret
}

// NoGlobalCache prevents reading/writing entities from/to memcache.
func (qry *Query) NoGlobalCache() (ret *Query) {
	ret = qry.clone()
	ret.opts.NoGlobalCache = true
	return
}

// Limit returns a derivative Query that has a limit on the number
// of results returned. A negative value means unlimited.
func (qry *Query) Limit(limit int) (ret *Query) {
	ret = qry.clone()
	if limit > 0 {
		//ret.log("LIMIT %v", limit)
	} else {
		limit = -1
		//ret.log("NO LIMIT")
	}
	ret.inner.Limit = limit
	return
}

// NoLimit returns a derivative Query that has no limit on the number
// of results returned.
func (qry *Query) NoLimit() (ret *Query) {
	return qry.Limit(-1)
}

// Ancestor returns a derivative Query with an ancestor filter.
// The ancestor should not be nil.
func (qry *Query) Ancestor(k *Key) (ret *Query) {
	ret = qry.clone()
	//ret.log("ANCESTOR '%v'", k.String())
	ret.inner.Ancestor = k.inner
	return
}

// Project returns a derivative Query that yields only the passed fields.
// It cannot be used in a keys-only query.
func (qry *Query) Project(fields ...string) (ret *Query) {
	ret = qry.clone()
	//ret.log("PROJECT '%v'", strings.Join(fields, "', '"))
	ret.inner.Projection = append([]string(nil), fields...)
	ret.inner.TypeOf = types.ProjectQuery
	return
}

// EventualConsistency returns a derivative query that returns eventually
// consistent results. It only has an effect on ancestor queries.
func (qry *Query) EventualConsistency() (ret *Query) {
	ret = qry.clone()
	//ret.log("EVENTUAL CONSISTENCY")
	ret.inner.Eventual = true
	return
}

// Start returns a derivative Query with the passed start point.
func (qry *Query) Start(c string) (ret *Query) {
	ret = qry.clone()
	//ret.log("START CURSOR")
	ret.inner.Start = c
	return
}

// End returns a derivative Query with the passed end point.
func (qry *Query) End(c string) (ret *Query) {
	ret = qry.clone()
	//ret.log("END CURSOR")
	ret.inner.End = c
	return
}

// Offset returns a derivative Query that has an offset of how many keys
// to skip over before returning results. A negative value is invalid.
func (qry *Query) Offset(off int) (ret *Query) {
	ret = qry.clone()
	//ret.log("OFFSET %v", off)
	ret.inner.Offset = off
	return
}

// OrderAsc returns a derivative Query with a field-based sort order, ascending.
// Orders are applied in the order they are added.
func (qry *Query) OrderAsc(s string) (ret *Query) {
	ret = qry.clone()
	//ret.log("ORDER ASC %v", s)
	ret.inner.Order = append(ret.inner.Order, types.Order{FieldName: s, Descending: false})
	return
}

// OrderDesc returns a derivative Query with a field-based sort order, descending.
// Orders are applied in the order they are added.
func (qry *Query) OrderDesc(s string) (ret *Query) {
	ret = qry.clone()
	//ret.log("ORDER DESC %v", s)
	ret.inner.Order = append(ret.inner.Order, types.Order{FieldName: s, Descending: true})
	return
}

// Distinct returns a derivative query that yields de-duplicated entities with
// respect to the set of projected fields. It is only used for projection
// queries.
func (qry *Query) Distinct() (ret *Query) {
	ret = qry.clone()
	ret.inner.Distinct = true
	return ret
}

// Filter returns a derivative Query with a field-based filter.
// The filterStr argument must be a field name followed by optional space,
// followed by an operator, one of ">", "<", ">=", "<=", or "=".
// Fields are compared against the provided value using the operator.
// Multiple filters are AND'ed together.
func (qry *Query) Filter(q string, val interface{}) (ret *Query) {
	ret = qry.clone()
	ret.inner.Filter = append(ret.inner.Filter, types.Filter{Filter: q, Value: val})
	//ret.log("FILTER '%v %v'", q, val)
	return
}

// GetCount returns the number of results for the query.
func (qry *Query) GetCount() (int, error) {
	//qry.log("COUNT")
	//qry.ctx.Infof(qry.getLog())

	return dsCount(qry.ctx, qry.inner)
}

// GetKeys executes the query as keys-only: No entities are retrieved, just their keys.
func (qry *Query) GetKeys() ([]*Key, string, error) {
	keysQry := qry.clone()
	//keysQry.log("KEYS-ONLY")
	keysQry.inner.TypeOf = types.KeysOnlyQuery

	it := keysQry.Run()

	keys, err := it.GetAll(nil)
	if err != nil {
		return nil, "", err
	}

	cursor, err := it.Cursor()
	return keys, cursor, err
}

// GetAll runs the query and writes the entities to the passed destination.
//
// Note that, if not manually disabled, queries for more than 1 item use
// a "hybrid query". This means that first a keys-only query is executed
// and then the keys are used to lookup the local and global cache as well
// as the datastore eventually. For a warm cache this usually is
// faster and cheaper than the regular query.
func (qry *Query) GetAll(dsts interface{}) ([]*Key, string, error) {
	useHybridQry := qry.inner.Limit != 1 && qry.inner.TypeOf == types.FullQuery && !qry.opts.NoGlobalCache
	if useHybridQry {
		return qry.getAllByHybrid(dsts)
	}

	it := qry.Run()
	keys, err := it.GetAll(dsts)
	if err != nil {
		return nil, "", err
	}
	cursor, err := it.Cursor()
	return keys, cursor, err
}

func (qry *Query) getAllByHybrid(dsts interface{}) ([]*Key, string, error) {
	keys, cursor, err := qry.GetKeys()
	if err == nil && len(keys) > 0 {
		keys, err = newLoader(qry.ctx, qry.kind).Keys(keys).GetAll(dsts)
	}
	return keys, cursor, err
}

// GetFirst executes the query and writes the result's first entity
// to the passed destination.
func (qry *Query) GetFirst(dst interface{}) (*Key, error) {
	return qry.Run().GetOne(dst)
}

// Run executes the query and returns an Iterator.
func (qry *Query) Run() *Iterator {
	//qry.ctx.Infof(qry.getLog())
	return newIterator(qry)
}

//func (qry *Query) log(s string, values ...interface{}) {
//	qry.log = append(qry.log, fmt.Sprintf(s, values...))
//}
//
//func (qry *Query) getLog() string {
//	return fmt.Sprintf("running query \"%v\"", strings.Join(qry.log, " | "))
//}
<|endoftext|>


<commit_before>package livestatus

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
)

// Query is a binding query instance.
type Query struct {
	table   string
	headers []string
	columns []string
	ls      *Livestatus
}

// Columns sets the names of the columns to retrieve when executing a query.
func (q *Query) Columns(names ...string) *Query {
	q.headers = append(q.headers, "Columns: "+strings.Join(names, " "))
	q.columns = names
	return q
}

// Filter sets a new filter to apply to the query.
func (q *Query) Filter(rule string) *Query {
	q.headers = append(q.headers, "Filter: "+rule)
	return q
}

// And combines the n last filters into a new filter using a `And` operation.
func (q *Query) And(n int) *Query {
	q.headers = append(q.headers, fmt.Sprintf("And: %d", n))
	return q
}

// Or combines the n last filters into a new filter using a `Or` operation.
func (q *Query) Or(n int) *Query {
	q.headers = append(q.headers, fmt.Sprintf("Or: %d", n))
	return q
}

// Negate negates the most recent filter.
func (q *Query) Negate() *Query {
	q.headers = append(q.headers, "Negate:")
	return q
}

// Exec executes the query.
func (q *Query) Exec() (*Response, error) {
	resp := &Response{}

	// Connect to socket
	conn, err := q.dial()
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	// Send command data
	conn.Write([]byte(q.buildCmd()))

	// Read response header
	data := make([]byte, 16)

	_, err = conn.Read(data)
	if err != nil {
		return nil, err
	}

	resp.Status, err = strconv.Atoi(string(data[:3]))
	if err != nil {
		return nil, err
	}

	// Receive response data
	buf := bytes.NewBuffer(nil)

	for {
		data = make([]byte, 1024)

		_, err = conn.Read(data)
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}

		buf.Write(bytes.TrimRight(data, "\x00"))
	}

	if buf.Len() == 0 {
		return resp, nil
	}

	// Parse received data for records
	resp.Records, err = q.parse(buf.Bytes())
	if err != nil {
		return nil, err
	}

	return resp, nil
}

func (q *Query) buildCmd() string {
	cmd := "GET " + q.table

	// Append headers if any
	if len(q.headers) > 0 {
		cmd += "\n" + strings.Join(q.headers, "\n")
	}

	// Set default headers
	cmd += "\nResponseHeader: fixed16"
	cmd += "\nOutputFormat: json"
	cmd += "\n\n"

	return cmd
}

func (q *Query) dial() (net.Conn, error) {
	return net.Dial(q.ls.network, q.ls.address)
}

func (q *Query) parse(data []byte) ([]Record, error) {
	var (
		records []Record
		rows    [][]interface{}
	)

	// Unmarshal received data
	if err := json.Unmarshal(data, &rows); err != nil {
		return nil, err
	} else if len(rows) < 2 {
		return records, nil
	}

	// Extract columns names from first row
	start := 0

	if len(q.columns) == 0 {
		q.columns = make([]string, len(rows[0]))
		for i, value := range rows[0] {
			q.columns[i] = value.(string)
		}

		start = 1
	}

	// Fill records maps
	for _, row := range rows[start:] {
		r := make(Record)
		for i, value := range row {
			r.set(q.columns[i], value)
		}
		records = append(records, r)
	}

	return records, nil
}

func newQuery(table string, ls *Livestatus) *Query {
	return &Query{
		table:   table,
		headers: make([]string, 0),
		ls:      ls,
	}
}
<commit_msg>Fix missing first row when columns are specified<commit_after>package livestatus

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
)

// Query is a binding query instance.
type Query struct {
	table   string
	headers []string
	columns []string
	ls      *Livestatus
}

// Columns sets the names of the columns to retrieve when executing a query.
func (q *Query) Columns(names ...string) *Query {
	q.headers = append(q.headers, "Columns: "+strings.Join(names, " "))
	q.columns = names
	return q
}

// Filter sets a new filter to apply to the query.
func (q *Query) Filter(rule string) *Query {
	q.headers = append(q.headers, "Filter: "+rule)
	return q
}

// And combines the n last filters into a new filter using a `And` operation.
func (q *Query) And(n int) *Query {
	q.headers = append(q.headers, fmt.Sprintf("And: %d", n))
	return q
}

// Or combines the n last filters into a new filter using a `Or` operation.
func (q *Query) Or(n int) *Query {
	q.headers = append(q.headers, fmt.Sprintf("Or: %d", n))
	return q
}

// Negate negates the most recent filter.
func (q *Query) Negate() *Query {
	q.headers = append(q.headers, "Negate:")
	return q
}

// Exec executes the query.
func (q *Query) Exec() (*Response, error) {
	resp := &Response{}

	// Connect to socket
	conn, err := q.dial()
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	// Send command data
	conn.Write([]byte(q.buildCmd()))

	// Read response header
	data := make([]byte, 16)

	_, err = conn.Read(data)
	if err != nil {
		return nil, err
	}

	resp.Status, err = strconv.Atoi(string(data[:3]))
	if err != nil {
		return nil, err
	}

	// Receive response data
	buf := bytes.NewBuffer(nil)

	for {
		data = make([]byte, 1024)

		_, err = conn.Read(data)
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}

		buf.Write(bytes.TrimRight(data, "\x00"))
	}

	if buf.Len() == 0 {
		return resp, nil
	}

	// Parse received data for records
	resp.Records, err = q.parse(buf.Bytes())
	if err != nil {
		return nil, err
	}

	return resp, nil
}

func (q *Query) buildCmd() string {
	cmd := "GET " + q.table

	// Append headers if any
	if len(q.headers) > 0 {
		cmd += "\n" + strings.Join(q.headers, "\n")
	}

	// Set default headers
	cmd += "\nResponseHeader: fixed16"
	cmd += "\nOutputFormat: json"
	cmd += "\n\n"

	return cmd
}

func (q *Query) dial() (net.Conn, error) {
	return net.Dial(q.ls.network, q.ls.address)
}

func (q *Query) parse(data []byte) ([]Record, error) {
	var (
		records []Record
		rows    [][]interface{}
	)

	// Unmarshal received data
	if err := json.Unmarshal(data, &rows); err != nil {
		return nil, err
	} else if len(q.columns) == 0 && len(rows) < 2 || len(q.columns) > 0 && len(rows) < 1 {
		return records, nil
	}

	// Extract columns names from first row
	start := 0

	if len(q.columns) == 0 {
		q.columns = make([]string, len(rows[0]))
		for i, value := range rows[0] {
			q.columns[i] = value.(string)
		}

		start = 1
	}

	// Fill records maps
	for _, row := range rows[start:] {
		r := make(Record)
		for i, value := range row {
			r.set(q.columns[i], value)
		}
		records = append(records, r)
	}

	return records, nil
}

func newQuery(table string, ls *Livestatus) *Query {
	return &Query{
		table:   table,
		headers: make([]string, 0),
		ls:      ls,
	}
}
<|endoftext|>


<commit_before>package arangolite

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

// Query represents an AQL query.
type Query struct {
	aql       string
	cache     bool
	batchSize int
}

// NewQuery returns a new Query object.
func NewQuery(aql string, params ...interface{}) *Query {
	aql = processAQLQuery(aql) // Process to remove eventual tabs/spaces used when indenting the query
	aql = fmt.Sprintf(aql, params...)
	aql = strings.Replace(aql, `"`, "'", -1) // Replace by single quotes so there is no conflict when serialised in JSON

	return &Query{aql: aql}
}

// Cache enables/disables the caching of the query.
// Unavailable prior to ArangoDB 2.7
func (q *Query) Cache(enable bool) *Query {
	q.cache = enable
	return q
}

// BatchSize sets the batch size of the query
func (q *Query) BatchSize(size int) *Query {
	q.batchSize = size
	return q
}

// Run runs the query synchronously and returns the JSON array of all elements
// of every batch returned by the database.
func (q *Query) Run(db *DB) ([]byte, error) {
	async, err := q.RunAsync(db)
	if err != nil {
		return nil, err
	}

	return db.syncResult(async), nil
}

// RunAsync runs the query asynchronously and returns an async Result object.
func (q *Query) RunAsync(db *DB) (*Result, error) {
	if db == nil {
		return nil, errors.New("nil database")
	}

	if len(q.aql) == 0 {
		return NewResult(nil), nil
	}

	c, err := db.runQuery("/_api/cursor", q)

	if err != nil {
		return nil, err
	}

	return NewResult(c), nil
}

func (q *Query) generate() []byte {
	type QueryFmt struct {
		Query     string `json:"query"`
		Cache     bool   `json:"cache"`
		BatchSize int    `json:"batchSize,omitempty"`
	}

	jsonQuery, _ := json.Marshal(&QueryFmt{Query: q.aql, Cache: q.cache, BatchSize: q.batchSize})

	return jsonQuery
}

func (q *Query) description() string {
	return "QUERY"
}

func processAQLQuery(query string) string {
	query = strings.Replace(query, "\n", " ", -1)
	query = strings.Replace(query, "\t", "", -1)

	split := strings.Split(query, " ")
	split2 := []string{}

	for _, s := range split {
		if len(s) == 0 {
			continue
		}
		split2 = append(split2, s)
	}

	query = strings.Join(split2, " ")

	return query
}
<commit_msg>Query cache field is now a pointer.<commit_after>package arangolite

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

// Query represents an AQL query.
type Query struct {
	aql       string
	cache     *bool
	batchSize int
}

// NewQuery returns a new Query object.
func NewQuery(aql string, params ...interface{}) *Query {
	aql = processAQLQuery(aql) // Process to remove eventual tabs/spaces used when indenting the query
	aql = fmt.Sprintf(aql, params...)
	aql = strings.Replace(aql, `"`, "'", -1) // Replace by single quotes so there is no conflict when serialised in JSON

	return &Query{aql: aql}
}

// Cache enables/disables the caching of the query.
// Unavailable prior to ArangoDB 2.7
func (q *Query) Cache(enable bool) *Query {
	q.cache = &enable
	return q
}

// BatchSize sets the batch size of the query
func (q *Query) BatchSize(size int) *Query {
	q.batchSize = size
	return q
}

// Run runs the query synchronously and returns the JSON array of all elements
// of every batch returned by the database.
func (q *Query) Run(db *DB) ([]byte, error) {
	async, err := q.RunAsync(db)
	if err != nil {
		return nil, err
	}

	return db.syncResult(async), nil
}

// RunAsync runs the query asynchronously and returns an async Result object.
func (q *Query) RunAsync(db *DB) (*Result, error) {
	if db == nil {
		return nil, errors.New("nil database")
	}

	if len(q.aql) == 0 {
		return NewResult(nil), nil
	}

	c, err := db.runQuery("/_api/cursor", q)

	if err != nil {
		return nil, err
	}

	return NewResult(c), nil
}

func (q *Query) generate() []byte {
	type QueryFmt struct {
		Query     string `json:"query"`
		Cache     *bool  `json:"cache,omitempty"`
		BatchSize int    `json:"batchSize,omitempty"`
	}

	jsonQuery, _ := json.Marshal(&QueryFmt{Query: q.aql, Cache: q.cache, BatchSize: q.batchSize})

	return jsonQuery
}

func (q *Query) description() string {
	return "QUERY"
}

func processAQLQuery(query string) string {
	query = strings.Replace(query, "\n", " ", -1)
	query = strings.Replace(query, "\t", "", -1)

	split := strings.Split(query, " ")
	split2 := []string{}

	for _, s := range split {
		if len(s) == 0 {
			continue
		}
		split2 = append(split2, s)
	}

	query = strings.Join(split2, " ")

	return query
}
<|endoftext|>