<commit_before>/** \copyright
 * Copyright (c) 2019, Mike Dunston
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are  permitted provided that the following conditions are met:
 *
 *  - Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 *  - Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 * \file Esp32HardwareCanAdapter.hxx
 *
 * ESP32 Hardware CAN adapter code. This utilizes the built in CAN controller
 * to translate the can_frame in OpenMRN to the ESP32 can_message_t used by the
 * ESP-IDF CAN controller code. The ESP32 will still require an external CAN
 * transceiver (MCP2551 or SN65HVD230 as example).
 *
 * @author Mike Dunston
 * @date 19 January 2019
 */

#ifndef _FREERTOS_DRIVERS_ESP32_ESP32HWCAN_HXX_
#define _FREERTOS_DRIVERS_ESP32_ESP32HWCAN_HXX_

#include "freertos_drivers/arduino/Can.hxx"
#include <driver/can.h>
#include <driver/gpio.h>
#include <esp_task_wdt.h>

/// ESP32 CAN bus status strings, used for periodic status reporting
static const char *ESP32_CAN_STATUS_STRINGS[] = {
    "STOPPED",               // CAN_STATE_STOPPED
    "RUNNING",               // CAN_STATE_RUNNING
    "OFF / RECOVERY NEEDED", // CAN_STATE_BUS_OFF
    "RECOVERY UNDERWAY"      // CAN_STATE_RECOVERING
};

class Esp32HardwareCan : public Can
{
public:
    /// Constructor.
    ///
    /// @param name is the name for the CAN adapter, this is currently not used.
    /// @param rxPin is the ESP32 pin that is connected to the external
    /// transceiver RX.
    /// @param txPin is the ESP32 pin that is connected to the external
    /// transceiver TX.
    Esp32HardwareCan(const char *name, gpio_num_t rxPin, gpio_num_t txPin, bool reportStats=true)
        : Can(name), reportStats_(reportStats)
    {
        // Configure the ESP32 CAN driver to use 125kbps.
        can_timing_config_t can_timing_config = CAN_TIMING_CONFIG_125KBITS();
        // By default we accept all CAN frames.
        can_filter_config_t can_filter_config = CAN_FILTER_CONFIG_ACCEPT_ALL();
        // Note: not using the CAN_GENERAL_CONFIG_DEFAULT macro due to a missing
        // cast for CAN_IO_UNUSED.
        can_general_config_t can_general_config = {.mode = CAN_MODE_NORMAL,
            .tx_io = txPin,
            .rx_io = rxPin,
            .clkout_io = (gpio_num_t)CAN_IO_UNUSED,
            .bus_off_io = (gpio_num_t)CAN_IO_UNUSED,
            .tx_queue_len = (uint32_t)config_can_tx_buffer_size() / 2,
            .rx_queue_len = (uint32_t)config_can_rx_buffer_size() / 2,
            .alerts_enabled = CAN_ALERT_NONE,
            .clkout_divider = 0};

        LOG(VERBOSE,
            "ESP32-CAN driver configured using RX: %d, TX: %d, TX-Q: %d, "
            "RX-Q: %d",
            can_general_config.rx_io, can_general_config.tx_io,
            can_general_config.rx_queue_len, can_general_config.tx_queue_len);

        ESP_ERROR_CHECK(can_driver_install(
            &can_general_config, &can_timing_config, &can_filter_config));

        xTaskCreatePinnedToCore(rx_task, "ESP32-CAN RX", OPENMRN_STACK_SIZE,
            this, RX_TASK_PRIORITY, &rxTaskHandle_, tskNO_AFFINITY);
        xTaskCreatePinnedToCore(tx_task, "ESP32-CAN TX", OPENMRN_STACK_SIZE,
            this, TX_TASK_PRIORITY, &txTaskHandle_, tskNO_AFFINITY);
    }

    ~Esp32HardwareCan()
    {
    }

    /// Enables the ESP32 CAN driver
    virtual void enable()
    {
        ESP_ERROR_CHECK(can_start());
        LOG(VERBOSE, "ESP32-CAN driver enabled");
    }

    /// Disables the ESP32 CAN driver
    virtual void disable()
    {
        ESP_ERROR_CHECK(can_stop());
        LOG(VERBOSE, "ESP32-CAN driver disabled");
    }

protected:
    /// function to try and transmit a message
    void tx_msg() override
    {
        // wake up the tx_task so it can consume any can_frames from txBuf.
        xTaskNotifyGive(txTaskHandle_);
    }

private:
    /// Default constructor.
    Esp32HardwareCan();

    /// Enables/Disables the periodic reporting of CAN bus statistics to the
    /// default serial stream.
    bool reportStats_;

    /// Handle for the tx_task that converts and transmits can_frame to the
    /// native can driver.
    TaskHandle_t txTaskHandle_;

    /// Handle for the rx_task that receives and converts the native can driver
    /// frames to can_frame.
    TaskHandle_t rxTaskHandle_;

    /// Interval at which to print the ESP32 CAN bus status.
    static constexpr TickType_t STATUS_PRINT_INTERVAL = pdMS_TO_TICKS(10000);

    /// Interval to wait between iterations when the bus is recovering, a
    /// transmit failure or there is nothing to transmit.
    static constexpr TickType_t TX_DEFAULT_DELAY = pdMS_TO_TICKS(250);

    /// Priority to use for the rx_task. This needs to be higher than the
    /// tx_task and lower than @ref OPENMRN_TASK_PRIORITY.
    static constexpr UBaseType_t RX_TASK_PRIORITY = ESP_TASK_TCPIP_PRIO - 1;

    /// Priority to use for the tx_task. This should be lower than
    /// @ref RX_TASK_PRIORITY and @ref OPENMRN_TASK_PRIORITY.
    static constexpr UBaseType_t TX_TASK_PRIORITY = ESP_TASK_TCPIP_PRIO - 2;

    /// Background task that takes care of the conversion of the @ref can_frame
    /// provided by the @ref txBuf into an ESP32 can_message_t which can be
    /// processed by the native CAN driver. This task also covers the periodic
    /// status reporting and BUS recovery when necessary.
    static void tx_task(void *can)
    {
        /// Get handle to our parent Esp32HardwareCan object to access the
        /// txBuf.
        Esp32HardwareCan *parent = reinterpret_cast<Esp32HardwareCan *>(can);

        // Add this task to the WDT
        esp_task_wdt_add(parent->txTaskHandle_);

        /// Tracks the last time that we displayed the CAN driver status.
        TickType_t next_status_display_tick_count = 0;

        while (true)
        {
            // Feed the watchdog so it doesn't reset the ESP32
            esp_task_wdt_reset();

            // periodic CAN driver monitoring and reporting, this takes care of
            // bus recovery when the CAN driver disables the bus due to error
            // conditions exceeding thresholds.
            can_status_info_t status;
            can_get_status_info(&status);
            auto current_tick_count = xTaskGetTickCount();
            if ((next_status_display_tick_count == 0 ||
                current_tick_count >= next_status_display_tick_count) &&
                parent->reportStats_)
            {
                next_status_display_tick_count =
                    current_tick_count + STATUS_PRINT_INTERVAL;
                LOG(INFO,
                    "ESP32-CAN: rx-q:%d, tx-q:%d, rx-err:%d, tx-err:%d, "
                    "arb-lost:%d, bus-err:%d, state: %s",
                    status.msgs_to_rx, status.msgs_to_tx,
                    status.rx_error_counter, status.tx_error_counter,
                    status.arb_lost_count, status.bus_error_count,
                    ESP32_CAN_STATUS_STRINGS[status.state]);
            }
            if (status.state == CAN_STATE_BUS_OFF)
            {
                // When the bus is OFF we need to initiate recovery, transmit is
                // not possible when in this state.
                LOG(WARNING, "ESP32-CAN: initiating recovery");
                can_initiate_recovery();
                continue;
            }
            else if (status.state == CAN_STATE_RECOVERING)
            {
                // when the bus is in recovery mode transmit is not possible.
                vTaskDelay(TX_DEFAULT_DELAY);
                continue;
            }

            // check txBuf for any message to transmit.
            unsigned count;
            struct can_frame *can_frame = nullptr;
            {
                AtomicHolder h(parent);
                count = parent->txBuf->data_read_pointer(&can_frame);
            }
            if (!count || !can_frame)
            {
                // tx Buf empty; wait for tx_msg to be called.
                ulTaskNotifyTake(pdTRUE, // clear on exit
                    TX_DEFAULT_DELAY);
                continue;
            }

            /// ESP32 native CAN driver frame
            can_message_t msg = {0};

            msg.flags = CAN_MSG_FLAG_NONE;
            msg.identifier = can_frame->can_id;
            msg.data_length_code = can_frame->can_dlc;
            for (int i = 0; i < can_frame->can_dlc; i++)
            {
                msg.data[i] = can_frame->data[i];
            }
            if (IS_CAN_FRAME_EFF(*can_frame))
            {
                msg.flags |= CAN_MSG_FLAG_EXTD;
            }
            if (IS_CAN_FRAME_RTR(*can_frame))
            {
                msg.flags |= CAN_MSG_FLAG_RTR;
            }

            // Pass the converted CAN frame to the native driver
            // for transmit, if the TX queue is full this will
            // return ESP_ERR_TIMEOUT which will result in the
            // the message being left in txBuf for the next iteration.
            // if this call returns ESP_OK we consider the frame as
            // transmitted by the driver and remove it from txBuf.
            esp_err_t tx_res = can_transmit(&msg, pdMS_TO_TICKS(100));
            if (tx_res == ESP_OK)
            {
                LOG(VERBOSE,
                    "ESP32-CAN-TX OK id:%08x, flags:%04x, dlc:%02d, "
                    "data:%02x%02x%02x%02x%02x%02x%02x%02x",
                    msg.identifier, msg.flags, msg.data_length_code,
                    msg.data[0], msg.data[1], msg.data[2], msg.data[3],
                    msg.data[4], msg.data[5], msg.data[6], msg.data[7]);
                AtomicHolder h(parent);
                parent->txBuf->consume(1);
                parent->txBuf->signal_condition();
            }
            else if (tx_res != ESP_ERR_TIMEOUT)
            {
                LOG(WARNING, "ESP32-CAN-TX: %s", esp_err_to_name(tx_res));
                vTaskDelay(TX_DEFAULT_DELAY);
            }
        } // loop on task
    }

    /// Background task that takes care of receiving can_message_t objects from
    /// the ESP32 native CAN driver, when they are available, converting them to
    /// a @ref can_frame and pushing them to the @ref rxBuf.
    static void rx_task(void *can)
    {
        /// Get handle to our parent Esp32HardwareCan object to access the rxBuf
        Esp32HardwareCan *parent = reinterpret_cast<Esp32HardwareCan *>(can);

        // Add this task to the WDT
        esp_task_wdt_add(parent->rxTaskHandle_);

        while (true)
        {
            // Feed the watchdog so it doesn't reset the ESP32
            esp_task_wdt_reset();

            /// ESP32 native CAN driver frame
            can_message_t msg = {0};
            if (can_receive(&msg, pdMS_TO_TICKS(250)) == ESP_OK)
            {
                // frame retrieved from the driver, convert to OpenMRN can_frame
                if (msg.flags & CAN_MSG_FLAG_DLC_NON_COMP)
                {
                    LOG(WARNING,
                        "ESP32-CAN-RX: received non-compliant CAN frame, frame "
                        "dropped!");
                }
                else
                {
                    LOG(VERBOSE,
                        "ESP32-CAN-RX id:%08x, flags:%04x, dlc:%02d, "
                        "data:%02x%02x%02x%02x%02x%02x%02x%02x",
                        msg.identifier, msg.flags, msg.data_length_code,
                        msg.data[0], msg.data[1], msg.data[2], msg.data[3],
                        msg.data[4], msg.data[5], msg.data[6], msg.data[7]);
                    AtomicHolder h(parent);
                    struct can_frame *can_frame = nullptr;
                    if (parent->rxBuf->data_write_pointer(&can_frame) &&
                        can_frame != nullptr)
                    {
                        LOG(VERBOSE, "ESP32-CAN-RX: converting to can_frame");
                        memset(can_frame, 0, sizeof(struct can_frame));
                        can_frame->can_id = msg.identifier;
                        can_frame->can_dlc = msg.data_length_code;
                        for (int i = 0; i < msg.data_length_code; i++)
                        {
                            can_frame->data[i] = msg.data[i];
                        }
                        if (msg.flags & CAN_MSG_FLAG_EXTD)
                        {
                            can_frame->can_eff = 1;
                        }
                        if (msg.flags & CAN_MSG_FLAG_RTR)
                        {
                            can_frame->can_rtr = 1;
                        }
                        parent->rxBuf->advance(1);
                        parent->rxBuf->signal_condition();
                    }
                    else
                    {
                        LOG(WARNING,
                            "ESP32-CAN-RX: buffer overrun, frame dropped!");
                        parent->overrunCount++;
                    }
                }
            }
        }
    }
    DISALLOW_COPY_AND_ASSIGN(Esp32HardwareCan);
};

#endif /* _FREERTOS_DRIVERS_ARDUINO_ESP32HWCAN_HXX_ */
<commit_msg>review comments<commit_after>/** \copyright
 * Copyright (c) 2019, Mike Dunston
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are  permitted provided that the following conditions are met:
 *
 *  - Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 *  - Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 * \file Esp32HardwareCanAdapter.hxx
 *
 * ESP32 Hardware CAN adapter code. This utilizes the built in CAN controller
 * to translate the can_frame in OpenMRN to the ESP32 can_message_t used by the
 * ESP-IDF CAN controller code. The ESP32 will still require an external CAN
 * transceiver (MCP2551 or SN65HVD230 as example).
 *
 * @author Mike Dunston
 * @date 19 January 2019
 */

#ifndef _FREERTOS_DRIVERS_ESP32_ESP32HWCAN_HXX_
#define _FREERTOS_DRIVERS_ESP32_ESP32HWCAN_HXX_

#include "freertos_drivers/arduino/Can.hxx"
#include <driver/can.h>
#include <driver/gpio.h>
#include <esp_task_wdt.h>

/// ESP32 CAN bus status strings, used for periodic status reporting
static const char *ESP32_CAN_STATUS_STRINGS[] = {
    "STOPPED",               // CAN_STATE_STOPPED
    "RUNNING",               // CAN_STATE_RUNNING
    "OFF / RECOVERY NEEDED", // CAN_STATE_BUS_OFF
    "RECOVERY UNDERWAY"      // CAN_STATE_RECOVERING
};

class Esp32HardwareCan : public Can
{
public:
    /// Constructor.
    ///
    /// @param name is the name for the CAN adapter, this is currently not used.
    /// @param rxPin is the ESP32 pin that is connected to the external
    /// transceiver RX.
    /// @param txPin is the ESP32 pin that is connected to the external
    /// transceiver TX.
    Esp32HardwareCan(const char *name, gpio_num_t rxPin, gpio_num_t txPin, bool reportStats=true)
        : Can(name), reportStats_(reportStats)
    {
        // Configure the ESP32 CAN driver to use 125kbps.
        can_timing_config_t can_timing_config = CAN_TIMING_CONFIG_125KBITS();
        // By default we accept all CAN frames.
        can_filter_config_t can_filter_config = CAN_FILTER_CONFIG_ACCEPT_ALL();
        // Note: not using the CAN_GENERAL_CONFIG_DEFAULT macro due to a missing
        // cast for CAN_IO_UNUSED.
        can_general_config_t can_general_config = {.mode = CAN_MODE_NORMAL,
            .tx_io = txPin,
            .rx_io = rxPin,
            .clkout_io = (gpio_num_t)CAN_IO_UNUSED,
            .bus_off_io = (gpio_num_t)CAN_IO_UNUSED,
            .tx_queue_len = (uint32_t)config_can_tx_buffer_size() / 2,
            .rx_queue_len = (uint32_t)config_can_rx_buffer_size() / 2,
            .alerts_enabled = CAN_ALERT_NONE,
            .clkout_divider = 0};

        LOG(VERBOSE,
            "ESP32-CAN driver configured using RX: %d, TX: %d, RX-Q: %d, "
            "TX-Q: %d",
            can_general_config.rx_io, can_general_config.tx_io,
            can_general_config.rx_queue_len, can_general_config.tx_queue_len);

        ESP_ERROR_CHECK(can_driver_install(
            &can_general_config, &can_timing_config, &can_filter_config));

        xTaskCreatePinnedToCore(rx_task, "ESP32-CAN RX", OPENMRN_STACK_SIZE,
            this, RX_TASK_PRIORITY, &rxTaskHandle_, tskNO_AFFINITY);
        xTaskCreatePinnedToCore(tx_task, "ESP32-CAN TX", OPENMRN_STACK_SIZE,
            this, TX_TASK_PRIORITY, &txTaskHandle_, tskNO_AFFINITY);
    }

    ~Esp32HardwareCan()
    {
    }

    /// Enables the ESP32 CAN driver
    virtual void enable()
    {
        ESP_ERROR_CHECK(can_start());
        LOG(VERBOSE, "ESP32-CAN driver enabled");
    }

    /// Disables the ESP32 CAN driver
    virtual void disable()
    {
        ESP_ERROR_CHECK(can_stop());
        LOG(VERBOSE, "ESP32-CAN driver disabled");
    }

protected:
    /// function to try and transmit a message
    void tx_msg() override
    {
        // wake up the tx_task so it can consume any can_frames from txBuf.
        xTaskNotifyGive(txTaskHandle_);
    }

private:
    /// Default constructor.
    Esp32HardwareCan();

    /// Enables/Disables the periodic reporting of CAN bus statistics to the
    /// default serial stream.
    bool reportStats_;

    /// Handle for the tx_task that converts and transmits can_frame to the
    /// native can driver.
    TaskHandle_t txTaskHandle_;

    /// Handle for the rx_task that receives and converts the native can driver
    /// frames to can_frame.
    TaskHandle_t rxTaskHandle_;

    /// Interval at which to print the ESP32 CAN bus status.
    static constexpr TickType_t STATUS_PRINT_INTERVAL = pdMS_TO_TICKS(10000);

    /// Interval to wait between iterations when the bus is recovering, a
    /// transmit failure or there is nothing to transmit.
    static constexpr TickType_t TX_DEFAULT_DELAY = pdMS_TO_TICKS(250);

    /// Priority to use for the rx_task. This needs to be higher than the
    /// tx_task and lower than @ref OPENMRN_TASK_PRIORITY.
    static constexpr UBaseType_t RX_TASK_PRIORITY = ESP_TASK_TCPIP_PRIO - 1;

    /// Priority to use for the tx_task. This should be lower than
    /// @ref RX_TASK_PRIORITY and @ref OPENMRN_TASK_PRIORITY.
    static constexpr UBaseType_t TX_TASK_PRIORITY = ESP_TASK_TCPIP_PRIO - 2;

    /// Background task that takes care of the conversion of the @ref can_frame
    /// provided by the @ref txBuf into an ESP32 can_message_t which can be
    /// processed by the native CAN driver. This task also covers the periodic
    /// status reporting and BUS recovery when necessary.
    static void tx_task(void *can)
    {
        /// Get handle to our parent Esp32HardwareCan object to access the
        /// txBuf.
        Esp32HardwareCan *parent = reinterpret_cast<Esp32HardwareCan *>(can);

        // Add this task to the WDT
        esp_task_wdt_add(parent->txTaskHandle_);

        /// Tracks the last time that we displayed the CAN driver status.
        TickType_t next_status_display_tick_count = 0;

        while (true)
        {
            // Feed the watchdog so it doesn't reset the ESP32
            esp_task_wdt_reset();

            // periodic CAN driver monitoring and reporting, this takes care of
            // bus recovery when the CAN driver disables the bus due to error
            // conditions exceeding thresholds.
            can_status_info_t status;
            can_get_status_info(&status);
            auto current_tick_count = xTaskGetTickCount();
            if ((next_status_display_tick_count == 0 ||
                current_tick_count >= next_status_display_tick_count) &&
                parent->reportStats_)
            {
                next_status_display_tick_count =
                    current_tick_count + STATUS_PRINT_INTERVAL;
                LOG(INFO,
                    "ESP32-CAN: rx-q:%d, tx-q:%d, rx-err:%d, tx-err:%d, "
                    "arb-lost:%d, bus-err:%d, state: %s",
                    status.msgs_to_rx, status.msgs_to_tx,
                    status.rx_error_counter, status.tx_error_counter,
                    status.arb_lost_count, status.bus_error_count,
                    ESP32_CAN_STATUS_STRINGS[status.state]);
            }
            if (status.state == CAN_STATE_BUS_OFF)
            {
                // When the bus is OFF we need to initiate recovery, transmit is
                // not possible when in this state.
                LOG(WARNING, "ESP32-CAN: initiating recovery");
                can_initiate_recovery();
                continue;
            }
            else if (status.state == CAN_STATE_RECOVERING)
            {
                // when the bus is in recovery mode transmit is not possible.
                vTaskDelay(TX_DEFAULT_DELAY);
                continue;
            }

            // check txBuf for any message to transmit.
            unsigned count;
            struct can_frame *can_frame = nullptr;
            {
                AtomicHolder h(parent);
                count = parent->txBuf->data_read_pointer(&can_frame);
            }
            if (!count || !can_frame)
            {
                // tx Buf empty; wait for tx_msg to be called.
                ulTaskNotifyTake(pdTRUE, // clear on exit
                    TX_DEFAULT_DELAY);
                continue;
            }

            /// ESP32 native CAN driver frame
            can_message_t msg = {0};

            msg.flags = CAN_MSG_FLAG_NONE;
            msg.identifier = can_frame->can_id;
            msg.data_length_code = can_frame->can_dlc;
            for (int i = 0; i < can_frame->can_dlc; i++)
            {
                msg.data[i] = can_frame->data[i];
            }
            if (IS_CAN_FRAME_EFF(*can_frame))
            {
                msg.flags |= CAN_MSG_FLAG_EXTD;
            }
            if (IS_CAN_FRAME_RTR(*can_frame))
            {
                msg.flags |= CAN_MSG_FLAG_RTR;
            }

            // Pass the converted CAN frame to the native driver
            // for transmit, if the TX queue is full this will
            // return ESP_ERR_TIMEOUT which will result in the
            // the message being left in txBuf for the next iteration.
            // if this call returns ESP_OK we consider the frame as
            // transmitted by the driver and remove it from txBuf.
            esp_err_t tx_res = can_transmit(&msg, pdMS_TO_TICKS(100));
            if (tx_res == ESP_OK)
            {
                LOG(VERBOSE,
                    "ESP32-CAN-TX OK id:%08x, flags:%04x, dlc:%02d, "
                    "data:%02x%02x%02x%02x%02x%02x%02x%02x",
                    msg.identifier, msg.flags, msg.data_length_code,
                    msg.data[0], msg.data[1], msg.data[2], msg.data[3],
                    msg.data[4], msg.data[5], msg.data[6], msg.data[7]);
                AtomicHolder h(parent);
                parent->txBuf->consume(1);
                parent->txBuf->signal_condition();
            }
            else if (tx_res != ESP_ERR_TIMEOUT)
            {
                LOG(WARNING, "ESP32-CAN-TX: %s", esp_err_to_name(tx_res));
                vTaskDelay(TX_DEFAULT_DELAY);
            }
        } // loop on task
    }

    /// Background task that takes care of receiving can_message_t objects from
    /// the ESP32 native CAN driver, when they are available, converting them to
    /// a @ref can_frame and pushing them to the @ref rxBuf.
    static void rx_task(void *can)
    {
        /// Get handle to our parent Esp32HardwareCan object to access the rxBuf
        Esp32HardwareCan *parent = reinterpret_cast<Esp32HardwareCan *>(can);

        // Add this task to the WDT
        esp_task_wdt_add(parent->rxTaskHandle_);

        while (true)
        {
            // Feed the watchdog so it doesn't reset the ESP32
            esp_task_wdt_reset();

            /// ESP32 native CAN driver frame
            can_message_t msg = {0};
            if (can_receive(&msg, pdMS_TO_TICKS(250)) != ESP_OK)
            {
                // native CAN driver did not give us a frame.
                continue;
            }
            // we have received a frame from the native CAN driver, verify if
            // it is a standard frame, if not we drop it.
            if (msg.flags & CAN_MSG_FLAG_DLC_NON_COMP)
            {
                LOG(WARNING,
                    "ESP32-CAN-RX: received non-compliant CAN frame, frame "
                    "dropped!");
                continue;
            }
            LOG(VERBOSE,
                "ESP32-CAN-RX id:%08x, flags:%04x, dlc:%02d, "
                "data:%02x%02x%02x%02x%02x%02x%02x%02x",
                msg.identifier, msg.flags, msg.data_length_code,
                msg.data[0], msg.data[1], msg.data[2], msg.data[3],
                msg.data[4], msg.data[5], msg.data[6], msg.data[7]);
            AtomicHolder h(parent);
            struct can_frame *can_frame = nullptr;
            // verify if we have space in the rxBuf, if not drop the frame and
            // record the overrun.
            if (!parent->rxBuf->data_write_pointer(&can_frame) ||
                can_frame == nullptr)
            {
                LOG(WARNING,
                    "ESP32-CAN-RX: buffer overrun, frame dropped!");
                parent->overrunCount++;
                continue;
            }
            // we have space in the rxBuf, start conversion
            LOG(VERBOSE, "ESP32-CAN-RX: converting to can_frame");
            memset(can_frame, 0, sizeof(struct can_frame));
            can_frame->can_id = msg.identifier;
            can_frame->can_dlc = msg.data_length_code;
            for (int i = 0; i < msg.data_length_code; i++)
            {
                can_frame->data[i] = msg.data[i];
            }
            if (msg.flags & CAN_MSG_FLAG_EXTD)
            {
                SET_CAN_FRAME_EFF(*can_frame);
            }
            if (msg.flags & CAN_MSG_FLAG_RTR)
            {
                SET_CAN_FRAME_RTR(*can_frame);
            }
            parent->rxBuf->advance(1);
            parent->rxBuf->signal_condition();
        }
    }
    DISALLOW_COPY_AND_ASSIGN(Esp32HardwareCan);
};

#endif /* _FREERTOS_DRIVERS_ARDUINO_ESP32HWCAN_HXX_ */
<|endoftext|>


<commit_before>/*****************************************************************************
 * utils.cpp: ActiveX control for VLC
 *****************************************************************************
 * Copyright (C) 2005 the VideoLAN team
 *
 * Authors: Damien Fouilleul <Damien.Fouilleul@laposte.net>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.
 *****************************************************************************/

#include "utils.h"

/*
** conversion facilities
*/

using namespace std;

char *CStrFromBSTR(UINT codePage, BSTR bstr)
{
    UINT len = SysStringLen(bstr);
    if( len > 0 )
    {
        size_t mblen = WideCharToMultiByte(codePage,
                0, bstr, len, NULL, 0, NULL, NULL);
        if( mblen > 0 )
        {
            char *buffer = (char *)CoTaskMemAlloc(mblen+1);
            ZeroMemory(buffer, mblen+1);
            if( WideCharToMultiByte(codePage, 0, bstr, len, buffer, mblen, NULL, NULL) )
            {
                buffer[mblen] = '\0';
                return buffer;
            }
        }
    }
    return NULL;
};

BSTR BSTRFromCStr(UINT codePage, LPCSTR s)
{
    int wideLen = MultiByteToWideChar(codePage, 0, s, -1, NULL, 0);
    if( wideLen > 0 )
    {
        WCHAR* wideStr = (WCHAR*)CoTaskMemAlloc(wideLen*sizeof(WCHAR));
        if( NULL != wideStr )
        {
            BSTR bstr;

            ZeroMemory(wideStr, wideLen*sizeof(WCHAR));
            MultiByteToWideChar(codePage, 0, s, -1, wideStr, wideLen);
            bstr = SysAllocStringLen(wideStr, wideLen);
            free(wideStr);

            return bstr;
        }
    }
    return NULL;
};

/*
**  properties
*/

HRESULT GetObjectProperty(LPUNKNOWN object, DISPID dispID, VARIANT& v)
{
    IDispatch *pDisp;
    HRESULT hr = object->QueryInterface(IID_IDispatch, (LPVOID *)&pDisp);
    if( SUCCEEDED(hr) )
    {
        DISPPARAMS dispparamsNoArgs = {NULL, NULL, 0, 0};
        VARIANT vres;
        VariantInit(&vres);
        hr = pDisp->Invoke(dispID, IID_NULL, LOCALE_USER_DEFAULT,
                DISPATCH_PROPERTYGET, &dispparamsNoArgs, &vres, NULL, NULL);
        if( SUCCEEDED(hr) )
        {
            if( V_VT(&v) != V_VT(&vres) )
            {
                hr = VariantChangeType(&v, &vres, 0, V_VT(&v));
                VariantClear(&vres);
            }
            else
            {
                v = vres;
            }
        }
        pDisp->Release();
    }
    return hr;
};

HDC CreateDevDC(DVTARGETDEVICE *ptd)
{
	HDC hdc=NULL;
	LPDEVNAMES lpDevNames;
	LPDEVMODE lpDevMode;
	LPTSTR lpszDriverName;
	LPTSTR lpszDeviceName;
	LPTSTR lpszPortName;

	if (ptd == NULL) {
		hdc = CreateDC(TEXT("DISPLAY"), NULL, NULL, NULL);
		goto errReturn;
	}

	lpDevNames = (LPDEVNAMES) ptd; // offset for size field

	if (ptd->tdExtDevmodeOffset == 0) {
		lpDevMode = NULL;
	}else{
		lpDevMode  = (LPDEVMODE) ((LPTSTR)ptd + ptd->tdExtDevmodeOffset);
	}

	lpszDriverName = (LPTSTR) lpDevNames + ptd->tdDriverNameOffset;
	lpszDeviceName = (LPTSTR) lpDevNames + ptd->tdDeviceNameOffset;
	lpszPortName   = (LPTSTR) lpDevNames + ptd->tdPortNameOffset;

	hdc = CreateDC(lpszDriverName, lpszDeviceName, lpszPortName, lpDevMode);

errReturn:
	return hdc;
};

#define HIMETRIC_PER_INCH 2540

void DPFromHimetric(HDC hdc, LPPOINT pt, int count)
{
    LONG lpX = GetDeviceCaps(hdc, LOGPIXELSX);
    LONG lpY = GetDeviceCaps(hdc, LOGPIXELSY);
    while( count-- )
    {
        pt->x = pt->x*lpX/HIMETRIC_PER_INCH;
        pt->y = pt->y*lpY/HIMETRIC_PER_INCH;
        ++pt;
    }
};

void HimetricFromDP(HDC hdc, LPPOINT pt, int count)
{
    LONG lpX = GetDeviceCaps(hdc, LOGPIXELSX);
    LONG lpY = GetDeviceCaps(hdc, LOGPIXELSY);
    while( count-- )
    {
        pt->x = pt->x*HIMETRIC_PER_INCH/lpX;
        pt->y = pt->y*HIMETRIC_PER_INCH/lpY;
        ++pt;
    }
};

<commit_msg>- minor fixes<commit_after>/*****************************************************************************
 * utils.cpp: ActiveX control for VLC
 *****************************************************************************
 * Copyright (C) 2005 the VideoLAN team
 *
 * Authors: Damien Fouilleul <Damien.Fouilleul@laposte.net>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.
 *****************************************************************************/

#include "utils.h"

/*
** conversion facilities
*/

using namespace std;

char *CStrFromBSTR(UINT codePage, BSTR bstr)
{
    UINT len = SysStringLen(bstr);
    if( len > 0 )
    {
        size_t mblen = WideCharToMultiByte(codePage,
                0, bstr, len, NULL, 0, NULL, NULL);
        if( mblen > 0 )
        {
            char *buffer = (char *)CoTaskMemAlloc(mblen+1);
            ZeroMemory(buffer, mblen+1);
            if( WideCharToMultiByte(codePage, 0, bstr, len, buffer, mblen, NULL, NULL) )
            {
                buffer[mblen] = '\0';
                return buffer;
            }
        }
    }
    return NULL;
};

BSTR BSTRFromCStr(UINT codePage, LPCSTR s)
{
    int wideLen = MultiByteToWideChar(codePage, 0, s, -1, NULL, 0);
    if( wideLen > 0 )
    {
        WCHAR* wideStr = (WCHAR*)CoTaskMemAlloc(wideLen*sizeof(WCHAR));
        if( NULL != wideStr )
        {
            BSTR bstr;

            ZeroMemory(wideStr, wideLen*sizeof(WCHAR));
            MultiByteToWideChar(codePage, 0, s, -1, wideStr, wideLen);
            bstr = SysAllocStringLen(wideStr, wideLen);
            CoTaskMemFree(wideStr);

            return bstr;
        }
    }
    return NULL;
};

/*
**  properties
*/

HRESULT GetObjectProperty(LPUNKNOWN object, DISPID dispID, VARIANT& v)
{
    IDispatch *pDisp;
    HRESULT hr = object->QueryInterface(IID_IDispatch, (LPVOID *)&pDisp);
    if( SUCCEEDED(hr) )
    {
        DISPPARAMS dispparamsNoArgs = {NULL, NULL, 0, 0};
        VARIANT vres;
        VariantInit(&vres);
        hr = pDisp->Invoke(dispID, IID_NULL, LOCALE_USER_DEFAULT,
                DISPATCH_PROPERTYGET, &dispparamsNoArgs, &vres, NULL, NULL);
        if( SUCCEEDED(hr) )
        {
            if( V_VT(&v) != V_VT(&vres) )
            {
                hr = VariantChangeType(&v, &vres, 0, V_VT(&v));
                VariantClear(&vres);
            }
            else
            {
                v = vres;
            }
        }
        pDisp->Release();
    }
    return hr;
};

HDC CreateDevDC(DVTARGETDEVICE *ptd)
{
	HDC hdc=NULL;
	if( NULL == ptd )
    {
		hdc = CreateDC(TEXT("DISPLAY"), NULL, NULL, NULL);
	}
    else
    {
        LPDEVNAMES lpDevNames;
        LPDEVMODE lpDevMode;
        LPTSTR lpszDriverName;
        LPTSTR lpszDeviceName;
        LPTSTR lpszPortName;

        lpDevNames = (LPDEVNAMES) ptd; // offset for size field

        if (ptd->tdExtDevmodeOffset == 0)
        {
            lpDevMode = NULL;
        }
        else
        {
            lpDevMode  = (LPDEVMODE) ((LPTSTR)ptd + ptd->tdExtDevmodeOffset);
        }

        lpszDriverName = (LPTSTR) lpDevNames + ptd->tdDriverNameOffset;
        lpszDeviceName = (LPTSTR) lpDevNames + ptd->tdDeviceNameOffset;
        lpszPortName   = (LPTSTR) lpDevNames + ptd->tdPortNameOffset;

        hdc = CreateDC(lpszDriverName, lpszDeviceName, lpszPortName, lpDevMode);
    }
	return hdc;
};

#define HIMETRIC_PER_INCH 2540

void DPFromHimetric(HDC hdc, LPPOINT pt, int count)
{
    LONG lpX = GetDeviceCaps(hdc, LOGPIXELSX);
    LONG lpY = GetDeviceCaps(hdc, LOGPIXELSY);
    while( count-- )
    {
        pt->x = pt->x*lpX/HIMETRIC_PER_INCH;
        pt->y = pt->y*lpY/HIMETRIC_PER_INCH;
        ++pt;
    }
};

void HimetricFromDP(HDC hdc, LPPOINT pt, int count)
{
    LONG lpX = GetDeviceCaps(hdc, LOGPIXELSX);
    LONG lpY = GetDeviceCaps(hdc, LOGPIXELSY);
    while( count-- )
    {
        pt->x = pt->x*HIMETRIC_PER_INCH/lpX;
        pt->y = pt->y*HIMETRIC_PER_INCH/lpY;
        ++pt;
    }
};

<|endoftext|>


<commit_before>// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/graph/function_impl.h"
#include "core/graph/graph_viewer.h"
#include "core/graph/model.h"
#include "onnx/shape_inference/implementation.h"

namespace onnxruntime {
// Auto inferred and generate an opschema for stand-alone functions
// TODO: revisit to see if we can eliminate typeconstraint step
void IOTypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto* onnx_func_proto_,
                            std::unique_ptr<ONNX_NAMESPACE::OpSchema>& op_schema_,
                            const std::unordered_map<std::string, int>& input_name_idx_map,
                            const std::unordered_map<std::string, int>& output_name_idx_map) {
  std::vector<std::pair<std::string, std::string>> input_types_list(onnx_func_proto_->input_size());
  std::vector<std::pair<std::string, std::string>> output_types_list(onnx_func_proto_->output_size());
  std::unordered_map<std::string, std::vector<std::string>> type_constraint_map;
  std::unordered_map<std::string, ONNX_NAMESPACE::AttributeProto_AttributeType> attribute_type_map;
  auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
  for (auto& node : onnx_func_proto_->node()) {
    const auto node_op_schema =
        schema_registry->GetSchema(node.op_type(), static_cast<int>(onnx_func_proto_->since_version()), node.domain());
    for (int i = 0; i < node.input_size(); ++i) {
      auto& in_name = node.input().Get(i);
      auto iter = input_name_idx_map.find(in_name);
      if (iter != input_name_idx_map.end()) {
        int idx = iter->second;
        const auto& p = node_op_schema->inputs().at(i);
        std::string type_str = p.GetTypeStr() + "in" + std::to_string(i);
        input_types_list[idx] = std::make_pair(in_name, type_str);
        if (!type_constraint_map.count(type_str)) {
          for (auto s : p.GetTypes()) {
            type_constraint_map[type_str].emplace_back(*s);
          }
        }
      }
    }
    for (int i = 0; i < node.output_size(); ++i) {
      auto& out_name = node.output().Get(i);
      auto iter = output_name_idx_map.find(out_name);
      if (iter != output_name_idx_map.end()) {
        int idx = iter->second;
        const auto& p = node_op_schema->outputs().at(i);
        std::string type_str = p.GetTypeStr() + "out" + std::to_string(i);
        output_types_list[idx] = std::make_pair(out_name, type_str);
        if (!type_constraint_map.count(type_str)) {
          for (auto s : p.GetTypes()) {
            type_constraint_map[type_str].emplace_back(*s);
          }
        }
      }
    }

    // If an subgraph node attribute has a specified
    // type attribute, we add its referenced attribute
    // into the op's schema
    for (auto& attr : node.attribute()) {
      if (attr.has_ref_attr_name() && attr.has_type())
        attribute_type_map[attr.ref_attr_name()] = attr.type();
    }
  }

  int i = 0;
  for (auto& input : input_types_list) {
    op_schema_->Input(i, input.first, "", input.second);
    ++i;
  }
  i = 0;
  for (auto& output : output_types_list) {
    op_schema_->Output(i, output.first, "", output.second);
    ++i;
  }

  for (auto& tc : type_constraint_map) {
    op_schema_->TypeConstraint(tc.first, tc.second, "");
  }

  for (auto& attribute_name : onnx_func_proto_->attribute()) {
    if (attribute_type_map.count(attribute_name))
      op_schema_->Attr(attribute_name, "", attribute_type_map[attribute_name], false);
  }
}

FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
                           std::unique_ptr<IndexedSubGraph> customized_func)
    : parent_graph_(&graph), onnx_func_proto_{nullptr} {
  customized_func_body_ = std::move(customized_func);

  // Construct body.
  body_ = std::make_unique<onnxruntime::Model>("fused_function_subgraph", false, onnxruntime::ModelMetaData(),
                                               IOnnxRuntimeOpSchemaRegistryList({graph.GetSchemaRegistry()}),
                                               graph.DomainToVersionMap());
  auto& sub_graph = body_->MainGraph();

  auto meta_def = customized_func_body_->GetMetaDef();
  op_schema_ = std::make_unique<ONNX_NAMESPACE::OpSchema>();
  op_schema_->SetName(meta_def->name);
  op_schema_->SetDomain(meta_def->domain);
  op_schema_->SetDoc(meta_def->doc_string);
  op_schema_->SinceVersion(meta_def->since_version);
  int i = 0;
  std::vector<const NodeArg*> sub_graph_inputs;
  sub_graph_inputs.resize(meta_def->inputs.size());
  for (auto& input : meta_def->inputs) {
    auto input_arg = parent_graph_->GetNodeArg(input);
    auto& sub_graph_input_arg = sub_graph.GetOrCreateNodeArg(input_arg->Name(), input_arg->TypeAsProto());
    sub_graph_inputs[i] = &sub_graph_input_arg;
    op_schema_->Input(i, input, "", *input_arg->Type());
    ++i;
  }
  i = 0;
  std::vector<const NodeArg*> sub_graph_outputs;
  sub_graph_outputs.resize(meta_def->outputs.size());
  for (auto& output : meta_def->outputs) {
    auto output_arg = parent_graph_->GetNodeArg(output);
    auto& sub_graph_output_arg = sub_graph.GetOrCreateNodeArg(output_arg->Name(), output_arg->TypeAsProto());
    sub_graph_outputs[i] = &sub_graph_output_arg;
    op_schema_->Output(i, output, "", *output_arg->Type());
    ++i;
  }
  op_schema_->Finalize();

  sub_graph.SetInputs(sub_graph_inputs);
  sub_graph.SetOutputs(sub_graph_outputs);
  //Add node and node args
  //TODO: for better performance, we could try to transfer the nodes in parent graph to sub-graph directly,
  //instead of create new nodes.
  for (auto& node_index : customized_func_body_->nodes) {
    auto node = parent_graph_->GetNode(node_index);
    std::vector<onnxruntime::NodeArg*> inputs;
    std::vector<onnxruntime::NodeArg*> outputs;
    for (auto input : node->InputDefs()) {
      auto& n_input = sub_graph.GetOrCreateNodeArg(input->Name(), input->TypeAsProto());
      inputs.push_back(&n_input);
    }
    for (auto output : node->OutputDefs()) {
      auto& n_output = sub_graph.GetOrCreateNodeArg(output->Name(), output->TypeAsProto());
      outputs.push_back(&n_output);
    }
    sub_graph.AddNode(node->Name(), node->OpType(), node->Description(), inputs, outputs, &node->GetAttributes(), node->Domain());
  }

  for (const auto& input : meta_def->inputs) {
    const ONNX_NAMESPACE::TensorProto* initializer = nullptr;
    if (graph.GetInitializedTensor(input, initializer)) {
      sub_graph.AddInitializedTensor(*initializer);
    }
  }

  //TODO: if we reuse the nodes in parent graph, maybe we don't need to resolve it.
  auto status = sub_graph.Resolve();
  ORT_ENFORCE(status.IsOK(), status.ErrorMessage());
}

FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
                           const onnxruntime::NodeIndex& node_index,
                           const ONNX_NAMESPACE::FunctionProto* onnx_func_proto)
    : parent_graph_(&graph) {
  onnx_func_proto_ = onnx_func_proto;
  auto node_in_parent_graph = parent_graph_->GetNode(node_index);
  op_schema_ = std::make_unique<ONNX_NAMESPACE::OpSchema>();
  op_schema_->SetName(onnx_func_proto_->name());
  op_schema_->SetDomain(onnx_func_proto_->node().Get(0).domain());
  op_schema_->SetDoc(onnx_func_proto_->doc_string());
  op_schema_->SinceVersion(static_cast<ONNX_NAMESPACE::OperatorSetVersion>(onnx_func_proto_->since_version()));
  std::unordered_map<std::string, int> input_name_idx_map;
  std::unordered_map<std::string, int> output_name_idx_map;
  for (int i = 0; i < onnx_func_proto_->input_size(); ++i) {
    input_name_idx_map[onnx_func_proto_->input().Get(i)] = i;
  }
  for (int i = 0; i < onnx_func_proto_->output_size(); ++i) {
    output_name_idx_map[onnx_func_proto_->output().Get(i)] = i;
  }

  auto cached_op_schema = node_in_parent_graph->Op();
  if (!cached_op_schema) {
    // Infer a op_schema for stand-alone functions.
    IOTypeConstraintHelper(onnx_func_proto_, this->op_schema_, input_name_idx_map, output_name_idx_map);
  } else {
    auto type_constraint_params = cached_op_schema->typeConstraintParams();
    for (auto& type_constraint_param : type_constraint_params) {
      op_schema_->TypeConstraint(
          type_constraint_param.type_param_str,
          type_constraint_param.allowed_type_strs,
          type_constraint_param.description);
    }
    int i = 0;
    for (auto& input : cached_op_schema->inputs()) {
      op_schema_->Input(i, input.GetName(), input.GetDescription(), input.GetTypeStr());
      ++i;
    }
    i = 0;
    for (auto& output : cached_op_schema->outputs()) {
      op_schema_->Output(i, output.GetName(), output.GetDescription(), output.GetTypeStr());
      ++i;
    }
    for (auto& attribute : cached_op_schema->attributes()) {
      op_schema_->Attr(attribute.second);
    }
  }

  if (!cached_op_schema || !cached_op_schema->has_type_and_shape_inference_function()) {
    op_schema_->TypeAndShapeInferenceFunction(
        [this](ONNX_NAMESPACE::InferenceContext& ctx) {
          auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
          const ONNX_NAMESPACE::FunctionProto* func_ptr = this->GetFuncProto();
          if (nullptr != func_ptr) {
            ONNX_NAMESPACE::shape_inference::InferShapeForFunctionNode(func_ptr, schema_registry, ctx);
          }
        });
  } else {
    op_schema_->TypeAndShapeInferenceFunction(cached_op_schema->GetTypeAndShapeInferenceFunction());
  }

  op_schema_->Finalize();
  //construct body
  std::unordered_map<std::string, int> domain_to_version;
  //TODO: set correct domain and version
  domain_to_version[onnxruntime::kOnnxDomain] = static_cast<int>(onnx_func_proto_->since_version());
  body_ = std::make_unique<onnxruntime::Model>(onnx_func_proto_->name(), false, onnxruntime::ModelMetaData(),
                                               IOnnxRuntimeOpSchemaRegistryList(), domain_to_version);
  auto& sub_graph = body_->MainGraph();
  // Add node and node args into subgraph
  // The subgraph preserved the input/output tensor names
  // in the parent graph for later inlining purpose
  auto attr_map = node_in_parent_graph->GetAttributes();
  for (auto& node : onnx_func_proto_->node()) {
    std::vector<onnxruntime::NodeArg*> inputs;
    std::vector<onnxruntime::NodeArg*> outputs;
    std::string uniq_identifier = node.name();
    if (!node.has_name()) {
      std::stringstream ss;
      ss << static_cast<const void*>(&node);
      uniq_identifier = ss.str();
    }

    for (int idx = 0; idx < node.input_size(); ++idx) {
      std::string tensor_name = node.input().Get(idx);
      auto iter = input_name_idx_map.find(tensor_name);
      if (iter != input_name_idx_map.end()) {
        // Preserving NodeArg and input/output names
        ONNX_NAMESPACE::NodeProto temp_node_proto;
        node_in_parent_graph->ToProto(temp_node_proto);
        const onnxruntime::NodeArg* node_arg = parent_graph_->GetNodeArg(temp_node_proto.input().Get(input_name_idx_map[tensor_name]));
        auto& n_input = sub_graph.GetOrCreateNodeArg(
            temp_node_proto.input().Get(iter->second), node_arg->TypeAsProto());
        inputs.push_back(&n_input);
      } else {
        auto& n_input = sub_graph.GetOrCreateNodeArg(
            tensor_name + "_" + std::to_string(node_index), nullptr);
        inputs.push_back(&n_input);
      }
    }
    for (int idx = 0; idx < node.output_size(); ++idx) {
      std::string tensor_name = node.output().Get(idx);
      auto iter = output_name_idx_map.find(tensor_name);
      if (iter != output_name_idx_map.end()) {
        // Preserving NodeArg and input/output names
        ONNX_NAMESPACE::NodeProto temp_node_proto;
        node_in_parent_graph->ToProto(temp_node_proto);
        const onnxruntime::NodeArg* node_arg = parent_graph_->GetNodeArg(temp_node_proto.output().Get(output_name_idx_map[tensor_name]));
        auto& n_output = sub_graph.GetOrCreateNodeArg(
            temp_node_proto.output().Get(iter->second), node_arg->TypeAsProto());
        outputs.push_back(&n_output);
      } else {
        auto& n_output = sub_graph.GetOrCreateNodeArg(
            tensor_name + "_" + std::to_string(node_index), nullptr);
        outputs.push_back(&n_output);
      }
    }

    onnxruntime::NodeAttributes new_attr_map;
    for (auto& attr : node.attribute()) {
      if (attr.has_ref_attr_name()) {
        if (attr_map.count(attr.ref_attr_name())) {
          new_attr_map[attr.name()] = attr_map[attr.ref_attr_name()];
        }
      } else {
        new_attr_map[attr.name()] = attr;
      }
    }
    sub_graph.AddNode(uniq_identifier + "_" + std::to_string(node_index), node.op_type(), node.doc_string(), inputs, outputs, &new_attr_map, node.domain());
  }
  auto status = sub_graph.Resolve();
  ORT_ENFORCE(status.IsOK(), "Resolve subgraph failed:", status.ErrorMessage());
}

FunctionImpl::~FunctionImpl() = default;

const ONNX_NAMESPACE::OpSchema& FunctionImpl::OpSchema() const {
  return *op_schema_;
}

const onnxruntime::Graph& FunctionImpl::Body() const {
  return body_->MainGraph();
}

const IndexedSubGraph& FunctionImpl::GetIndexedSubGraph() const {
  return *customized_func_body_;
}

const ONNX_NAMESPACE::FunctionProto* FunctionImpl::GetFuncProto() const {
  return onnx_func_proto_;
}

std::unique_ptr<Function> MakeFunction(const onnxruntime::Graph& graph,
                                       std::unique_ptr<IndexedSubGraph> customized_func) {
  return std::make_unique<FunctionImpl>(graph, std::move(customized_func));
}
}  // namespace onnxruntime
<commit_msg>Temp fix for a crash in fused graph (#1488)<commit_after>// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/graph/function_impl.h"
#include "core/graph/graph_viewer.h"
#include "core/graph/model.h"
#include "onnx/shape_inference/implementation.h"

namespace onnxruntime {
// Auto inferred and generate an opschema for stand-alone functions
// TODO: revisit to see if we can eliminate typeconstraint step
void IOTypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto* onnx_func_proto_,
                            std::unique_ptr<ONNX_NAMESPACE::OpSchema>& op_schema_,
                            const std::unordered_map<std::string, int>& input_name_idx_map,
                            const std::unordered_map<std::string, int>& output_name_idx_map) {
  std::vector<std::pair<std::string, std::string>> input_types_list(onnx_func_proto_->input_size());
  std::vector<std::pair<std::string, std::string>> output_types_list(onnx_func_proto_->output_size());
  std::unordered_map<std::string, std::vector<std::string>> type_constraint_map;
  std::unordered_map<std::string, ONNX_NAMESPACE::AttributeProto_AttributeType> attribute_type_map;
  auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
  for (auto& node : onnx_func_proto_->node()) {
    const auto node_op_schema =
        schema_registry->GetSchema(node.op_type(), static_cast<int>(onnx_func_proto_->since_version()), node.domain());
    for (int i = 0; i < node.input_size(); ++i) {
      auto& in_name = node.input().Get(i);
      auto iter = input_name_idx_map.find(in_name);
      if (iter != input_name_idx_map.end()) {
        int idx = iter->second;
        const auto& p = node_op_schema->inputs().at(i);
        std::string type_str = p.GetTypeStr() + "in" + std::to_string(i);
        input_types_list[idx] = std::make_pair(in_name, type_str);
        if (!type_constraint_map.count(type_str)) {
          for (auto s : p.GetTypes()) {
            type_constraint_map[type_str].emplace_back(*s);
          }
        }
      }
    }
    for (int i = 0; i < node.output_size(); ++i) {
      auto& out_name = node.output().Get(i);
      auto iter = output_name_idx_map.find(out_name);
      if (iter != output_name_idx_map.end()) {
        int idx = iter->second;
        const auto& p = node_op_schema->outputs().at(i);
        std::string type_str = p.GetTypeStr() + "out" + std::to_string(i);
        output_types_list[idx] = std::make_pair(out_name, type_str);
        if (!type_constraint_map.count(type_str)) {
          for (auto s : p.GetTypes()) {
            type_constraint_map[type_str].emplace_back(*s);
          }
        }
      }
    }

    // If an subgraph node attribute has a specified
    // type attribute, we add its referenced attribute
    // into the op's schema
    for (auto& attr : node.attribute()) {
      if (attr.has_ref_attr_name() && attr.has_type())
        attribute_type_map[attr.ref_attr_name()] = attr.type();
    }
  }

  int i = 0;
  for (auto& input : input_types_list) {
    op_schema_->Input(i, input.first, "", input.second);
    ++i;
  }
  i = 0;
  for (auto& output : output_types_list) {
    op_schema_->Output(i, output.first, "", output.second);
    ++i;
  }

  for (auto& tc : type_constraint_map) {
    op_schema_->TypeConstraint(tc.first, tc.second, "");
  }

  for (auto& attribute_name : onnx_func_proto_->attribute()) {
    if (attribute_type_map.count(attribute_name))
      op_schema_->Attr(attribute_name, "", attribute_type_map[attribute_name], false);
  }
}

FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
                           std::unique_ptr<IndexedSubGraph> customized_func)
    : parent_graph_(&graph), onnx_func_proto_{nullptr} {
  customized_func_body_ = std::move(customized_func);

  // Construct body.
  body_ = std::make_unique<onnxruntime::Model>("fused_function_subgraph", false, onnxruntime::ModelMetaData(),
                                               IOnnxRuntimeOpSchemaRegistryList({graph.GetSchemaRegistry()}),
                                               graph.DomainToVersionMap());
  auto& sub_graph = body_->MainGraph();

  auto meta_def = customized_func_body_->GetMetaDef();
  op_schema_ = std::make_unique<ONNX_NAMESPACE::OpSchema>();
  op_schema_->SetName(meta_def->name);
  op_schema_->SetDomain(meta_def->domain);
  op_schema_->SetDoc(meta_def->doc_string);
  op_schema_->SinceVersion(meta_def->since_version);
  int i = 0;
  std::vector<const NodeArg*> sub_graph_inputs;
  sub_graph_inputs.resize(meta_def->inputs.size());
  for (auto& input : meta_def->inputs) {
    auto input_arg = parent_graph_->GetNodeArg(input);
    auto& sub_graph_input_arg = sub_graph.GetOrCreateNodeArg(input_arg->Name(), input_arg->TypeAsProto());
    sub_graph_inputs[i] = &sub_graph_input_arg;
    ORT_ENFORCE(input_arg->Type() != nullptr);
    op_schema_->Input(i, input, "", *input_arg->Type());
    ++i;
  }
  i = 0;
  std::vector<const NodeArg*> sub_graph_outputs;
  sub_graph_outputs.resize(meta_def->outputs.size());
  for (auto& output : meta_def->outputs) {
    auto output_arg = parent_graph_->GetNodeArg(output);
    auto& sub_graph_output_arg = sub_graph.GetOrCreateNodeArg(output_arg->Name(), output_arg->TypeAsProto());
    sub_graph_outputs[i] = &sub_graph_output_arg;
    op_schema_->Output(i, output, "", *output_arg->Type());
    ++i;
  }
  op_schema_->Finalize();

  sub_graph.SetInputs(sub_graph_inputs);
  sub_graph.SetOutputs(sub_graph_outputs);
  //Add node and node args
  //TODO: for better performance, we could try to transfer the nodes in parent graph to sub-graph directly,
  //instead of create new nodes.
  for (auto& node_index : customized_func_body_->nodes) {
    auto node = parent_graph_->GetNode(node_index);
    std::vector<onnxruntime::NodeArg*> inputs;
    std::vector<onnxruntime::NodeArg*> outputs;
    for (auto input : node->InputDefs()) {
      auto& n_input = sub_graph.GetOrCreateNodeArg(input->Name(), input->TypeAsProto());
      inputs.push_back(&n_input);
    }
    for (auto output : node->OutputDefs()) {
      auto& n_output = sub_graph.GetOrCreateNodeArg(output->Name(), output->TypeAsProto());
      outputs.push_back(&n_output);
    }
    sub_graph.AddNode(node->Name(), node->OpType(), node->Description(), inputs, outputs, &node->GetAttributes(), node->Domain());
  }

  for (const auto& input : meta_def->inputs) {
    const ONNX_NAMESPACE::TensorProto* initializer = nullptr;
    if (graph.GetInitializedTensor(input, initializer)) {
      sub_graph.AddInitializedTensor(*initializer);
    }
  }

  //TODO: if we reuse the nodes in parent graph, maybe we don't need to resolve it.
  auto status = sub_graph.Resolve();
  ORT_ENFORCE(status.IsOK(), status.ErrorMessage());
}

FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
                           const onnxruntime::NodeIndex& node_index,
                           const ONNX_NAMESPACE::FunctionProto* onnx_func_proto)
    : parent_graph_(&graph) {
  onnx_func_proto_ = onnx_func_proto;
  auto node_in_parent_graph = parent_graph_->GetNode(node_index);
  op_schema_ = std::make_unique<ONNX_NAMESPACE::OpSchema>();
  op_schema_->SetName(onnx_func_proto_->name());
  op_schema_->SetDomain(onnx_func_proto_->node().Get(0).domain());
  op_schema_->SetDoc(onnx_func_proto_->doc_string());
  op_schema_->SinceVersion(static_cast<ONNX_NAMESPACE::OperatorSetVersion>(onnx_func_proto_->since_version()));
  std::unordered_map<std::string, int> input_name_idx_map;
  std::unordered_map<std::string, int> output_name_idx_map;
  for (int i = 0; i < onnx_func_proto_->input_size(); ++i) {
    input_name_idx_map[onnx_func_proto_->input().Get(i)] = i;
  }
  for (int i = 0; i < onnx_func_proto_->output_size(); ++i) {
    output_name_idx_map[onnx_func_proto_->output().Get(i)] = i;
  }

  auto cached_op_schema = node_in_parent_graph->Op();
  if (!cached_op_schema) {
    // Infer a op_schema for stand-alone functions.
    IOTypeConstraintHelper(onnx_func_proto_, this->op_schema_, input_name_idx_map, output_name_idx_map);
  } else {
    auto type_constraint_params = cached_op_schema->typeConstraintParams();
    for (auto& type_constraint_param : type_constraint_params) {
      op_schema_->TypeConstraint(
          type_constraint_param.type_param_str,
          type_constraint_param.allowed_type_strs,
          type_constraint_param.description);
    }
    int i = 0;
    for (auto& input : cached_op_schema->inputs()) {
      op_schema_->Input(i, input.GetName(), input.GetDescription(), input.GetTypeStr());
      ++i;
    }
    i = 0;
    for (auto& output : cached_op_schema->outputs()) {
      op_schema_->Output(i, output.GetName(), output.GetDescription(), output.GetTypeStr());
      ++i;
    }
    for (auto& attribute : cached_op_schema->attributes()) {
      op_schema_->Attr(attribute.second);
    }
  }

  if (!cached_op_schema || !cached_op_schema->has_type_and_shape_inference_function()) {
    op_schema_->TypeAndShapeInferenceFunction(
        [this](ONNX_NAMESPACE::InferenceContext& ctx) {
          auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
          const ONNX_NAMESPACE::FunctionProto* func_ptr = this->GetFuncProto();
          if (nullptr != func_ptr) {
            ONNX_NAMESPACE::shape_inference::InferShapeForFunctionNode(func_ptr, schema_registry, ctx);
          }
        });
  } else {
    op_schema_->TypeAndShapeInferenceFunction(cached_op_schema->GetTypeAndShapeInferenceFunction());
  }

  op_schema_->Finalize();
  //construct body
  std::unordered_map<std::string, int> domain_to_version;
  //TODO: set correct domain and version
  domain_to_version[onnxruntime::kOnnxDomain] = static_cast<int>(onnx_func_proto_->since_version());
  body_ = std::make_unique<onnxruntime::Model>(onnx_func_proto_->name(), false, onnxruntime::ModelMetaData(),
                                               IOnnxRuntimeOpSchemaRegistryList(), domain_to_version);
  auto& sub_graph = body_->MainGraph();
  // Add node and node args into subgraph
  // The subgraph preserved the input/output tensor names
  // in the parent graph for later inlining purpose
  auto attr_map = node_in_parent_graph->GetAttributes();
  for (auto& node : onnx_func_proto_->node()) {
    std::vector<onnxruntime::NodeArg*> inputs;
    std::vector<onnxruntime::NodeArg*> outputs;
    std::string uniq_identifier = node.name();
    if (!node.has_name()) {
      std::stringstream ss;
      ss << static_cast<const void*>(&node);
      uniq_identifier = ss.str();
    }

    for (int idx = 0; idx < node.input_size(); ++idx) {
      std::string tensor_name = node.input().Get(idx);
      auto iter = input_name_idx_map.find(tensor_name);
      if (iter != input_name_idx_map.end()) {
        // Preserving NodeArg and input/output names
        ONNX_NAMESPACE::NodeProto temp_node_proto;
        node_in_parent_graph->ToProto(temp_node_proto);
        const onnxruntime::NodeArg* node_arg = parent_graph_->GetNodeArg(temp_node_proto.input().Get(input_name_idx_map[tensor_name]));
        auto& n_input = sub_graph.GetOrCreateNodeArg(
            temp_node_proto.input().Get(iter->second), node_arg->TypeAsProto());
        inputs.push_back(&n_input);
      } else {
        auto& n_input = sub_graph.GetOrCreateNodeArg(
            tensor_name + "_" + std::to_string(node_index), nullptr);
        inputs.push_back(&n_input);
      }
    }
    for (int idx = 0; idx < node.output_size(); ++idx) {
      std::string tensor_name = node.output().Get(idx);
      auto iter = output_name_idx_map.find(tensor_name);
      if (iter != output_name_idx_map.end()) {
        // Preserving NodeArg and input/output names
        ONNX_NAMESPACE::NodeProto temp_node_proto;
        node_in_parent_graph->ToProto(temp_node_proto);
        const onnxruntime::NodeArg* node_arg = parent_graph_->GetNodeArg(temp_node_proto.output().Get(output_name_idx_map[tensor_name]));
        auto& n_output = sub_graph.GetOrCreateNodeArg(
            temp_node_proto.output().Get(iter->second), node_arg->TypeAsProto());
        outputs.push_back(&n_output);
      } else {
        auto& n_output = sub_graph.GetOrCreateNodeArg(
            tensor_name + "_" + std::to_string(node_index), nullptr);
        outputs.push_back(&n_output);
      }
    }

    onnxruntime::NodeAttributes new_attr_map;
    for (auto& attr : node.attribute()) {
      if (attr.has_ref_attr_name()) {
        if (attr_map.count(attr.ref_attr_name())) {
          new_attr_map[attr.name()] = attr_map[attr.ref_attr_name()];
        }
      } else {
        new_attr_map[attr.name()] = attr;
      }
    }
    sub_graph.AddNode(uniq_identifier + "_" + std::to_string(node_index), node.op_type(), node.doc_string(), inputs, outputs, &new_attr_map, node.domain());
  }
  auto status = sub_graph.Resolve();
  ORT_ENFORCE(status.IsOK(), "Resolve subgraph failed:", status.ErrorMessage());
}

FunctionImpl::~FunctionImpl() = default;

const ONNX_NAMESPACE::OpSchema& FunctionImpl::OpSchema() const {
  return *op_schema_;
}

const onnxruntime::Graph& FunctionImpl::Body() const {
  return body_->MainGraph();
}

const IndexedSubGraph& FunctionImpl::GetIndexedSubGraph() const {
  return *customized_func_body_;
}

const ONNX_NAMESPACE::FunctionProto* FunctionImpl::GetFuncProto() const {
  return onnx_func_proto_;
}

std::unique_ptr<Function> MakeFunction(const onnxruntime::Graph& graph,
                                       std::unique_ptr<IndexedSubGraph> customized_func) {
  return std::make_unique<FunctionImpl>(graph, std::move(customized_func));
}
}  // namespace onnxruntime
<|endoftext|>


<commit_before>#include <b9/interpreter.hpp>
#include <b9/jit.hpp>
#include <b9/loader.hpp>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>

/// B9run's usage string. Printed when run with -help.
static const char* usage =
    "Usage: b9run [<option>...] [--] <module> [<main>]\n"
    "   Or: b9run -help\n"
    "Jit Options:\n"
    "  -jit:         Enable the jit\n"
    "  -directcall:  make direct jit to jit calls\n"
    "  -passparam:   Pass arguments in CPU registers\n"
    "  -lazyvmstate: Only update the VM state as needed\n"
    "Run Options:\n"
    "  -loop <n>:    Run the program <n> times\n"
    "  -inline <n>:  Enable inlining\n"
    "  -debug:       Enable debug code\n"
    "  -verbose:     Run with verbose printing\n"
    "  -help:        Print this help message";

/// The b9run program's global configuration.
struct RunConfig {
  b9::Config b9;
  const char* moduleName = "";
  const char* mainFunction = "b9main";
  std::size_t loopCount = 1;
  bool verbose = false;
  std::vector<b9::StackElement> usrArgs;
};

/// Parse CLI arguments and set up the config.
static bool parseArguments(RunConfig& cfg, const int argc, char* argv[]) {
  std::size_t i = 1;

  for (; i < argc; i++) {
    const char* arg = argv[i];

    if (strcmp(arg, "-help") == 0) {
      std::cout << usage << std::endl;
      exit(EXIT_SUCCESS);
    } else if (strcmp(arg, "-loop") == 0) {
      cfg.loopCount = atoi(argv[++i]);
    } else if (strcmp(arg, "-inline") == 0) {
      cfg.b9.maxInlineDepth = atoi(argv[++i]);
    } else if (strcmp(arg, "-verbose") == 0) {
      cfg.verbose = true;
      cfg.b9.verbose = true;
    } else if (strcmp(arg, "-debug") == 0) {
      cfg.b9.debug = true;
    } else if (strcmp(arg, "-function") == 0) {
      cfg.mainFunction = argv[++i];
    } else if (strcmp(arg, "-jit") == 0) {
      cfg.b9.jit = true;
    } else if (strcmp(arg, "-directcall") == 0) {
      cfg.b9.directCall = true;
    } else if (strcmp(arg, "-passparam") == 0) {
      cfg.b9.passParam = true;
    } else if (strcmp(arg, "-lazyvmstate") == 0) {
      cfg.b9.lazyVmState = true;
    } else if (strcmp(arg, "--") == 0) {
      i++;
      break;
    } else if (strcmp(arg, "-") == 0) {
      std::cerr << "Unrecognized option: " << arg << std::endl;
      return false;
    } else {
      break;
    }
  }

  // check for user defined module
  if (i < argc) {
    cfg.moduleName = argv[i++];
  } else {
    std::cerr << "No module name given to b9run" << std::endl;
    return false;
  }

  // check for user defined arguments
  for (; i < argc; i++) {
    cfg.usrArgs.push_back(std::atoi(argv[i]));
  }

  return true;
}

static void run(const RunConfig& cfg) {
  b9::VirtualMachine vm{cfg.b9};

  b9::DlLoader loader{true};
  auto module = loader.loadModule(cfg.moduleName);

  if (module->functions.size() == 0) {
    throw b9::DlException{"Empty module"};
  }

  vm.load(module);

  if (cfg.b9.jit) vm.generateAllCode();

  size_t functionIndex = module->findFunction(cfg.mainFunction);
  if (cfg.loopCount == 1) {
    b9::StackElement returnVal = vm.run(functionIndex, cfg.usrArgs);
    std::cout << std::endl
              << cfg.mainFunction << " returned: " << returnVal << std::endl;
  } else {
    b9::StackElement returnVal;
    std::cout << "\nRunning " << cfg.mainFunction << " " << cfg.loopCount
              << " times:" << std::endl;
    for (std::size_t i = 1; i <= cfg.loopCount; i++) {
      returnVal = vm.run(functionIndex, cfg.usrArgs);
      std::cout << "Return value of iteration " << i << ": " << returnVal
                << std::endl;
    }
  }
}

int main(int argc, char* argv[]) {
  RunConfig cfg;

  if (!parseArguments(cfg, argc, argv)) {
    std::cerr << usage << std::endl;
    exit(EXIT_FAILURE);
  }

  try {
    run(cfg);
  } catch (const b9::DlException& e) {
    std::cerr << "Failed to load module: " << e.what() << std::endl;
    exit(EXIT_FAILURE);
  } catch (const b9::FunctionNotFoundException& e) {
    std::cerr << "Failed to find function: " << e.what() << std::endl;
    exit(EXIT_FAILURE);
  } catch (const b9::BadFunctionCallException& e) {
    std::cerr << "Failed to call function " << e.what() << std::endl;
    exit(EXIT_FAILURE);
  } catch (const b9::CompilationException& e) {
    std::cerr << "Failed to compile function: " << e.what() << std::endl;
    exit(EXIT_FAILURE);
  }

  exit(EXIT_SUCCESS);
}
<commit_msg>Do not check that the module is empty<commit_after>#include <b9/interpreter.hpp>
#include <b9/jit.hpp>
#include <b9/loader.hpp>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>

/// B9run's usage string. Printed when run with -help.
static const char* usage =
    "Usage: b9run [<option>...] [--] <module> [<main>]\n"
    "   Or: b9run -help\n"
    "Jit Options:\n"
    "  -jit:         Enable the jit\n"
    "  -directcall:  make direct jit to jit calls\n"
    "  -passparam:   Pass arguments in CPU registers\n"
    "  -lazyvmstate: Only update the VM state as needed\n"
    "Run Options:\n"
    "  -loop <n>:    Run the program <n> times\n"
    "  -inline <n>:  Enable inlining\n"
    "  -debug:       Enable debug code\n"
    "  -verbose:     Run with verbose printing\n"
    "  -help:        Print this help message";

/// The b9run program's global configuration.
struct RunConfig {
  b9::Config b9;
  const char* moduleName = "";
  const char* mainFunction = "b9main";
  std::size_t loopCount = 1;
  bool verbose = false;
  std::vector<b9::StackElement> usrArgs;
};

/// Parse CLI arguments and set up the config.
static bool parseArguments(RunConfig& cfg, const int argc, char* argv[]) {
  std::size_t i = 1;

  for (; i < argc; i++) {
    const char* arg = argv[i];

    if (strcmp(arg, "-help") == 0) {
      std::cout << usage << std::endl;
      exit(EXIT_SUCCESS);
    } else if (strcmp(arg, "-loop") == 0) {
      cfg.loopCount = atoi(argv[++i]);
    } else if (strcmp(arg, "-inline") == 0) {
      cfg.b9.maxInlineDepth = atoi(argv[++i]);
    } else if (strcmp(arg, "-verbose") == 0) {
      cfg.verbose = true;
      cfg.b9.verbose = true;
    } else if (strcmp(arg, "-debug") == 0) {
      cfg.b9.debug = true;
    } else if (strcmp(arg, "-function") == 0) {
      cfg.mainFunction = argv[++i];
    } else if (strcmp(arg, "-jit") == 0) {
      cfg.b9.jit = true;
    } else if (strcmp(arg, "-directcall") == 0) {
      cfg.b9.directCall = true;
    } else if (strcmp(arg, "-passparam") == 0) {
      cfg.b9.passParam = true;
    } else if (strcmp(arg, "-lazyvmstate") == 0) {
      cfg.b9.lazyVmState = true;
    } else if (strcmp(arg, "--") == 0) {
      i++;
      break;
    } else if (strcmp(arg, "-") == 0) {
      std::cerr << "Unrecognized option: " << arg << std::endl;
      return false;
    } else {
      break;
    }
  }

  // check for user defined module
  if (i < argc) {
    cfg.moduleName = argv[i++];
  } else {
    std::cerr << "No module name given to b9run" << std::endl;
    return false;
  }

  // check for user defined arguments
  for (; i < argc; i++) {
    cfg.usrArgs.push_back(std::atoi(argv[i]));
  }

  return true;
}

static void run(const RunConfig& cfg) {
  b9::VirtualMachine vm{cfg.b9};
  b9::DlLoader loader{true};

  auto module = loader.loadModule(cfg.moduleName);
  vm.load(module);

  if (cfg.b9.jit) vm.generateAllCode();

  size_t functionIndex = module->findFunction(cfg.mainFunction);
  if (cfg.loopCount == 1) {
    b9::StackElement returnVal = vm.run(functionIndex, cfg.usrArgs);
    std::cout << std::endl
              << cfg.mainFunction << " returned: " << returnVal << std::endl;
  } else {
    b9::StackElement returnVal;
    std::cout << "\nRunning " << cfg.mainFunction << " " << cfg.loopCount
              << " times:" << std::endl;
    for (std::size_t i = 1; i <= cfg.loopCount; i++) {
      returnVal = vm.run(functionIndex, cfg.usrArgs);
      std::cout << "Return value of iteration " << i << ": " << returnVal
                << std::endl;
    }
  }
}

int main(int argc, char* argv[]) {
  RunConfig cfg;

  if (!parseArguments(cfg, argc, argv)) {
    std::cerr << usage << std::endl;
    exit(EXIT_FAILURE);
  }

  try {
    run(cfg);
  } catch (const b9::DlException& e) {
    std::cerr << "Failed to load module: " << e.what() << std::endl;
    exit(EXIT_FAILURE);
  } catch (const b9::FunctionNotFoundException& e) {
    std::cerr << "Failed to find function: " << e.what() << std::endl;
    exit(EXIT_FAILURE);
  } catch (const b9::BadFunctionCallException& e) {
    std::cerr << "Failed to call function " << e.what() << std::endl;
    exit(EXIT_FAILURE);
  } catch (const b9::CompilationException& e) {
    std::cerr << "Failed to compile function: " << e.what() << std::endl;
    exit(EXIT_FAILURE);
  }

  exit(EXIT_SUCCESS);
}
<|endoftext|>


<commit_before>#include "WorldDragWidget.h"
#include <KrisLibrary/GLdraw/drawextra.h>
using namespace GLDraw;
using namespace Klampt;

WorldDragWidget::WorldDragWidget(WorldModel* _world)
  :world(_world),active(true),robotsActive(true),objectsActive(true),terrainsActive(false),
   highlightColor(1,1,1,0.3),lineColor(1,0.5,0),lineWidth(5.0),dragging(false),hoverID(-1),highlightID(-1)
{}

void WorldDragWidget::Set(WorldModel* _world)
{
  world = _world;
}

void WorldDragWidget::Enable(bool _active)
{
  active = _active;
}

void WorldDragWidget::SetHighlight(bool value)
{
  hasHighlight = value;
  if(hasHighlight && hoverID >= 0) {
    //update the object's color
    originalFaceColor = world->GetAppearance(hoverID)->faceColor;
    world->GetAppearance(hoverID)->faceColor.blend(originalFaceColor,highlightColor,highlightColor.rgba[3]);
    highlightID = hoverID;
  }
  else if(!hasHighlight && highlightID >= 0) {
    //restore the object's color
    world->GetAppearance(highlightID)->faceColor = originalFaceColor;
  }
}

bool WorldDragWidget::Hover(int x,int y,Camera::Viewport& viewport,double& distance)
{
  if(!active) return false;
  Ray3D r;
  viewport.getClickSource(x,y,r.source);
  viewport.getClickVector(x,y,r.direction);
  hoverID = -1;
  distance = Inf;
  if(robotsActive) {
    int body;
    Vector3 localpt;
    RobotModel* rob = world->RayCastRobot(r,body,localpt);
    if(rob) {
      hoverPt = localpt;
      int index = -1;
      for(size_t i=0;i<world->robots.size();i++)
        if(rob == world->robots[i].get()) { index=(int)i; break; }
      hoverID = world->RobotLinkID(index,body);
      auto geom = world->GetGeometry(hoverID);
      Vector3 worldpt = geom->GetTransform()*localpt;
      distance = worldpt.distance(r.source);
      dragPt = worldpt;
    }
  }
  if(objectsActive) {
    Vector3 localpt;
    RigidObjectModel* obj = world->RayCastObject(r,localpt);
    if(obj) {
      Vector3 worldpt = obj->T*localpt;
      Real d=worldpt.distance(r.source);
      if(d < distance) {
        distance = d;
        hoverPt = localpt;
        int index = -1;
        for(size_t i=0;i<world->rigidObjects.size();i++)
          if(obj == world->rigidObjects[i].get()) { index=(int)i; break; }
        hoverID = world->RigidObjectID(index);
        dragPt = worldpt;
      }
    }
  }
  hoverDistance = distance;
  if(hoverID >= 0)  {
    return true;
  }
  return false;
}

bool WorldDragWidget::BeginDrag(int x,int y,Camera::Viewport& viewport,double& distance)
{
  if(!Hover(x,y,viewport,distance)) return false;
  dragging = true;
  return true;
}

void WorldDragWidget::Drag(int dx,int dy,Camera::Viewport& viewport)
{
  Vector3 v;
  viewport.getMovementVectorAtDistance(dx,dy,hoverDistance,v);
  dragPt += v;
}

void WorldDragWidget::EndDrag()
{
  dragging = false;
}

void WorldDragWidget::DrawGL(Camera::Viewport& viewport)
{
  if(hoverID < 0) return;
  requestRedraw = false;
  if(hasFocus) {
    auto geom = world->GetGeometry(hoverID);
    glDisable(GL_LIGHTING);
    lineColor.setCurrentGL();
    glLineWidth(lineWidth);
    glBegin(GL_LINES);
    glVertex3v(geom->GetTransform()*hoverPt);
    glVertex3v(dragPt);
    glEnd();
    glLineWidth(1);
  }
}
<commit_msg>Fixed problem with force application mode turning items transprent<commit_after>#include "WorldDragWidget.h"
#include <KrisLibrary/GLdraw/drawextra.h>
using namespace GLDraw;
using namespace Klampt;

WorldDragWidget::WorldDragWidget(WorldModel* _world)
  :world(_world),active(true),robotsActive(true),objectsActive(true),terrainsActive(false),
   highlightColor(1,1,1,1.0),lineColor(1,0.5,0),lineWidth(5.0),dragging(false),hoverID(-1),highlightID(-1)
{}

void WorldDragWidget::Set(WorldModel* _world)
{
  world = _world;
}

void WorldDragWidget::Enable(bool _active)
{
  active = _active;
}

void WorldDragWidget::SetHighlight(bool value)
{
  hasHighlight = value;
  if(hasHighlight && hoverID >= 0) {
    //update the object's color
    originalFaceColor = world->GetAppearance(hoverID)->tintColor;
    world->GetAppearance(hoverID)->SetTintColor(highlightColor,0.3);
    highlightID = hoverID;
  }
  else if(!hasHighlight && highlightID >= 0) {
    //restore the object's color
    world->GetAppearance(highlightID)->tintColor = originalFaceColor;
    world->GetAppearance(highlightID)->tintStrength = 0.0;
  }
}

bool WorldDragWidget::Hover(int x,int y,Camera::Viewport& viewport,double& distance)
{
  if(!active) return false;
  Ray3D r;
  viewport.getClickSource(x,y,r.source);
  viewport.getClickVector(x,y,r.direction);
  hoverID = -1;
  distance = Inf;
  if(robotsActive) {
    int body;
    Vector3 localpt;
    RobotModel* rob = world->RayCastRobot(r,body,localpt);
    if(rob) {
      hoverPt = localpt;
      int index = -1;
      for(size_t i=0;i<world->robots.size();i++)
        if(rob == world->robots[i].get()) { index=(int)i; break; }
      hoverID = world->RobotLinkID(index,body);
      auto geom = world->GetGeometry(hoverID);
      Vector3 worldpt = geom->GetTransform()*localpt;
      distance = worldpt.distance(r.source);
      dragPt = worldpt;
    }
  }
  if(objectsActive) {
    Vector3 localpt;
    RigidObjectModel* obj = world->RayCastObject(r,localpt);
    if(obj) {
      Vector3 worldpt = obj->T*localpt;
      Real d=worldpt.distance(r.source);
      if(d < distance) {
        distance = d;
        hoverPt = localpt;
        int index = -1;
        for(size_t i=0;i<world->rigidObjects.size();i++)
          if(obj == world->rigidObjects[i].get()) { index=(int)i; break; }
        hoverID = world->RigidObjectID(index);
        dragPt = worldpt;
      }
    }
  }
  hoverDistance = distance;
  if(hoverID >= 0)  {
    return true;
  }
  return false;
}

bool WorldDragWidget::BeginDrag(int x,int y,Camera::Viewport& viewport,double& distance)
{
  if(!Hover(x,y,viewport,distance)) return false;
  dragging = true;
  return true;
}

void WorldDragWidget::Drag(int dx,int dy,Camera::Viewport& viewport)
{
  Vector3 v;
  viewport.getMovementVectorAtDistance(dx,dy,hoverDistance,v);
  dragPt += v;
}

void WorldDragWidget::EndDrag()
{
  dragging = false;
}

void WorldDragWidget::DrawGL(Camera::Viewport& viewport)
{
  if(hoverID < 0) return;
  requestRedraw = false;
  if(hasFocus) {
    auto geom = world->GetGeometry(hoverID);
    glDisable(GL_LIGHTING);
    lineColor.setCurrentGL();
    glLineWidth(lineWidth);
    glBegin(GL_LINES);
    glVertex3v(geom->GetTransform()*hoverPt);
    glVertex3v(dragPt);
    glEnd();
    glLineWidth(1);
  }
}
<|endoftext|>


<commit_before>/**
 * This file is part of the CernVM File System.
 */

#include "gtest/gtest.h"

#include <unistd.h>

#include <cstdio>

#include "../../cvmfs/compression.h"
#include "../../cvmfs/download.h"
#include "../../cvmfs/hash.h"
#include "../../cvmfs/prng.h"
#include "../../cvmfs/sink.h"
#include "../../cvmfs/statistics.h"
#include "../../cvmfs/util.h"

using namespace std;  // NOLINT

namespace download {

class T_Download : public ::testing::Test {
 protected:
  virtual void SetUp() {
    download_mgr.Init(8, false, /* use_system_proxy */
        &statistics);
    ffoo = CreateTempFile("./cvmfs_ut_download", 0600, "w+", &foo_path);
    assert(ffoo);
    foo_url = "file://" + foo_path;
  }

  virtual ~T_Download() {
    download_mgr.Fini();
    fclose(ffoo);
    unlink(foo_path.c_str());
  }

  perf::Statistics statistics;
  DownloadManager download_mgr;
  FILE *ffoo;
  string foo_path;
  string foo_url;
};

class TestSink : public cvmfs::Sink {
 public:
  TestSink() {
    FILE *f = CreateTempFile("./cvmfs_ut_download", 0600, "w+", &path);
    assert(f);
    fd = dup(fileno(f));
    assert(f >= 0);
    fclose(f);
  }

  virtual int64_t Write(const void *buf, uint64_t size) {
    return write(fd, buf, size);
  }

  virtual int Reset() {
    int retval = ftruncate(fd, 0);
    assert(retval == 0);
    return 0;
  }

  ~TestSink() {
    close(fd);
    unlink(path.c_str());
  }

  int fd;
  string path;
};

//------------------------------------------------------------------------------

// A placeholder test for future unit testing of the download module
TEST_F(T_Download, File) {
  string dest_path;
  FILE *fdest = CreateTempFile("./cvmfs_ut_download", 0600, "w+", &dest_path);
  ASSERT_TRUE(fdest != NULL);
  UnlinkGuard unlink_guard(dest_path);

  JobInfo info(&foo_url, false /* compressed */, false /* probe hosts */,
               fdest,  NULL);
  download_mgr.Fetch(&info);
  EXPECT_EQ(info.error_code, kFailOk);
  fclose(fdest);
}

TEST_F(T_Download, LocalFile2Mem) {
  string dest_path;
  FILE *fdest = CreateTempFile("./cvmfs_ut_download", 0600, "w+", &dest_path);
  ASSERT_TRUE(fdest != NULL);
  UnlinkGuard unlink_guard(dest_path);
  char buf = '1';
  fwrite(&buf, 1, 1, fdest);
  fclose(fdest);

  string url = "file://" + dest_path;
  JobInfo info(&url, false /* compressed */, false /* probe hosts */, NULL);
  download_mgr.Fetch(&info);
  ASSERT_EQ(info.error_code, kFailOk);
  ASSERT_EQ(info.destination_mem.size, 1U);
  EXPECT_EQ(info.destination_mem.data[0], '1');
}

TEST_F(T_Download, LocalFile2Sink) {
  string dest_path;
  FILE *fdest = CreateTempFile("./cvmfs_ut_download", 0600, "w+", &dest_path);
  ASSERT_TRUE(fdest != NULL);
  UnlinkGuard unlink_guard(dest_path);
  char buf = '1';
  fwrite(&buf, 1, 1, fdest);
  fflush(fdest);

  TestSink test_sink;
  string url = "file://" + dest_path;
  JobInfo info(&url, false /* compressed */, false /* probe hosts */,
               &test_sink, NULL /* expected hash */);
  download_mgr.Fetch(&info);
  EXPECT_EQ(info.error_code, kFailOk);
  EXPECT_EQ(1, pread(test_sink.fd, &buf, 1, 0));
  EXPECT_EQ('1', buf);

  rewind(fdest);
  Prng prng;
  prng.InitLocaltime();
  unsigned N = 16*1024;
  unsigned size = N*sizeof(uint32_t);
  uint32_t rnd_buf[N];  // 64kB
  for (unsigned i = 0; i < N; ++i)
    rnd_buf[i] = prng.Next(2147483647);
  shash::Any checksum(shash::kMd5);
  EXPECT_TRUE(
    zlib::CompressMem2File(reinterpret_cast<const unsigned char *>(rnd_buf),
                           size, fdest, &checksum));
  fclose(fdest);

  TestSink test_sink2;
  JobInfo info2(&url, true /* compressed */, false /* probe hosts */,
                &test_sink2, &checksum /* expected hash */);
  download_mgr.Fetch(&info2);
  EXPECT_EQ(info2.error_code, kFailOk);
  EXPECT_EQ(size, GetFileSize(test_sink2.path));

  uint32_t validation[N];
  EXPECT_EQ(static_cast<int>(size),
    pread(test_sink2.fd, &validation, size, 0));
  EXPECT_EQ(0, memcmp(validation, rnd_buf, size));
}

TEST_F(T_Download, StripDirect) {
  string cleaned = "FALSE";
  EXPECT_FALSE(download_mgr.StripDirect("", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("DIRECT", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("DIRECT;DIRECT", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("DIRECT;DIRECT|DIRECT", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("DIRECT;DIRECT|", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect(";", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect(";||;;;|||", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_FALSE(download_mgr.StripDirect("A|B", &cleaned));
  EXPECT_EQ("A|B", cleaned);
  EXPECT_FALSE(download_mgr.StripDirect("A|B;C|D;E|F|G", &cleaned));
  EXPECT_EQ("A|B;C|D;E|F|G", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("A|DIRECT;C|D;E|F;DIRECT", &cleaned));
  EXPECT_EQ("A;C|D;E|F", cleaned);
}

TEST_F(T_Download, ValidateGeoReply) {
  vector<uint64_t> geo_order;
  EXPECT_FALSE(download_mgr.ValidateGeoReply("", geo_order.size(), &geo_order));

  geo_order.push_back(0);
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("a", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("1,1", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("1,3", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("2,3", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("2", geo_order.size(), &geo_order));
  EXPECT_TRUE(
    download_mgr.ValidateGeoReply("1", geo_order.size(), &geo_order));
  EXPECT_EQ(geo_order.size(), 1U);
  EXPECT_EQ(geo_order[0], 0U);

  geo_order.push_back(0);
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply(",", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("2,", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("1", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("3,2,1", geo_order.size(), &geo_order));
  EXPECT_TRUE(
    download_mgr.ValidateGeoReply("2,1", geo_order.size(), &geo_order));
  EXPECT_EQ(geo_order.size(), 2U);
  EXPECT_EQ(geo_order[0], 1U);
  EXPECT_EQ(geo_order[1], 0U);

  EXPECT_TRUE(
    download_mgr.ValidateGeoReply("2,1\n", geo_order.size(), &geo_order));
  EXPECT_EQ(geo_order.size(), 2U);
  EXPECT_EQ(geo_order[0], 1U);
  EXPECT_EQ(geo_order[1], 0U);

  geo_order.push_back(0);
  geo_order.push_back(0);
  EXPECT_TRUE(
    download_mgr.ValidateGeoReply("4,3,1,2\n", geo_order.size(), &geo_order));
  EXPECT_EQ(geo_order.size(), 4U);
  EXPECT_EQ(geo_order[0], 3U);
  EXPECT_EQ(geo_order[1], 2U);
  EXPECT_EQ(geo_order[2], 0U);
  EXPECT_EQ(geo_order[3], 1U);
}

}  // namespace download
<commit_msg>FIX: T_Download depends on absolute paths<commit_after>/**
 * This file is part of the CernVM File System.
 */

#include "gtest/gtest.h"

#include <unistd.h>

#include <cstdio>

#include "../../cvmfs/compression.h"
#include "../../cvmfs/download.h"
#include "../../cvmfs/hash.h"
#include "../../cvmfs/prng.h"
#include "../../cvmfs/sink.h"
#include "../../cvmfs/statistics.h"
#include "../../cvmfs/util.h"

using namespace std;  // NOLINT

namespace download {

class T_Download : public ::testing::Test {
 protected:
  virtual void SetUp() {
    download_mgr.Init(8, false, /* use_system_proxy */
        &statistics);
    ffoo = CreateTemporaryFile(&foo_path);
    assert(ffoo);
    foo_url = "file://" + foo_path;
  }

  virtual ~T_Download() {
    download_mgr.Fini();
    fclose(ffoo);
    unlink(foo_path.c_str());
  }

  FILE *CreateTemporaryFile(std::string *path) const {
    return CreateTempFile(GetCurrentWorkingDirectory() + "/cvmfs_ut_download",
                          0600, "w+", path);
  }

  perf::Statistics statistics;
  DownloadManager download_mgr;
  FILE *ffoo;
  string foo_path;
  string foo_url;
};

class TestSink : public cvmfs::Sink {
 public:
  TestSink() {
    FILE *f = CreateTempFile("./cvmfs_ut_download", 0600, "w+", &path);
    assert(f);
    fd = dup(fileno(f));
    assert(f >= 0);
    fclose(f);
  }

  virtual int64_t Write(const void *buf, uint64_t size) {
    return write(fd, buf, size);
  }

  virtual int Reset() {
    int retval = ftruncate(fd, 0);
    assert(retval == 0);
    return 0;
  }

  ~TestSink() {
    close(fd);
    unlink(path.c_str());
  }

  int fd;
  string path;
};

//------------------------------------------------------------------------------

// A placeholder test for future unit testing of the download module
TEST_F(T_Download, File) {
  string dest_path;
  FILE *fdest = CreateTemporaryFile(&dest_path);
  ASSERT_TRUE(fdest != NULL);
  UnlinkGuard unlink_guard(dest_path);

  JobInfo info(&foo_url, false /* compressed */, false /* probe hosts */,
               fdest,  NULL);
  download_mgr.Fetch(&info);
  EXPECT_EQ(info.error_code, kFailOk);
  fclose(fdest);
}

TEST_F(T_Download, LocalFile2Mem) {
  string dest_path;
  FILE *fdest = CreateTemporaryFile(&dest_path);
  ASSERT_TRUE(fdest != NULL);
  UnlinkGuard unlink_guard(dest_path);
  char buf = '1';
  fwrite(&buf, 1, 1, fdest);
  fclose(fdest);

  string url = "file://" + dest_path;
  JobInfo info(&url, false /* compressed */, false /* probe hosts */, NULL);
  download_mgr.Fetch(&info);
  ASSERT_EQ(info.error_code, kFailOk);
  ASSERT_EQ(info.destination_mem.size, 1U);
  EXPECT_EQ(info.destination_mem.data[0], '1');
}

TEST_F(T_Download, LocalFile2Sink) {
  string dest_path;
  FILE *fdest = CreateTemporaryFile(&dest_path);
  ASSERT_TRUE(fdest != NULL);
  UnlinkGuard unlink_guard(dest_path);
  char buf = '1';
  fwrite(&buf, 1, 1, fdest);
  fflush(fdest);

  TestSink test_sink;
  string url = "file://" + dest_path;
  JobInfo info(&url, false /* compressed */, false /* probe hosts */,
               &test_sink, NULL /* expected hash */);
  download_mgr.Fetch(&info);
  EXPECT_EQ(info.error_code, kFailOk);
  EXPECT_EQ(1, pread(test_sink.fd, &buf, 1, 0));
  EXPECT_EQ('1', buf);

  rewind(fdest);
  Prng prng;
  prng.InitLocaltime();
  unsigned N = 16*1024;
  unsigned size = N*sizeof(uint32_t);
  uint32_t rnd_buf[N];  // 64kB
  for (unsigned i = 0; i < N; ++i)
    rnd_buf[i] = prng.Next(2147483647);
  shash::Any checksum(shash::kMd5);
  EXPECT_TRUE(
    zlib::CompressMem2File(reinterpret_cast<const unsigned char *>(rnd_buf),
                           size, fdest, &checksum));
  fclose(fdest);

  TestSink test_sink2;
  JobInfo info2(&url, true /* compressed */, false /* probe hosts */,
                &test_sink2, &checksum /* expected hash */);
  download_mgr.Fetch(&info2);
  EXPECT_EQ(info2.error_code, kFailOk);
  EXPECT_EQ(size, GetFileSize(test_sink2.path));

  uint32_t validation[N];
  EXPECT_EQ(static_cast<int>(size),
    pread(test_sink2.fd, &validation, size, 0));
  EXPECT_EQ(0, memcmp(validation, rnd_buf, size));
}

TEST_F(T_Download, StripDirect) {
  string cleaned = "FALSE";
  EXPECT_FALSE(download_mgr.StripDirect("", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("DIRECT", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("DIRECT;DIRECT", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("DIRECT;DIRECT|DIRECT", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("DIRECT;DIRECT|", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect(";", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect(";||;;;|||", &cleaned));
  EXPECT_EQ("", cleaned);
  EXPECT_FALSE(download_mgr.StripDirect("A|B", &cleaned));
  EXPECT_EQ("A|B", cleaned);
  EXPECT_FALSE(download_mgr.StripDirect("A|B;C|D;E|F|G", &cleaned));
  EXPECT_EQ("A|B;C|D;E|F|G", cleaned);
  EXPECT_TRUE(download_mgr.StripDirect("A|DIRECT;C|D;E|F;DIRECT", &cleaned));
  EXPECT_EQ("A;C|D;E|F", cleaned);
}

TEST_F(T_Download, ValidateGeoReply) {
  vector<uint64_t> geo_order;
  EXPECT_FALSE(download_mgr.ValidateGeoReply("", geo_order.size(), &geo_order));

  geo_order.push_back(0);
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("a", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("1,1", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("1,3", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("2,3", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("2", geo_order.size(), &geo_order));
  EXPECT_TRUE(
    download_mgr.ValidateGeoReply("1", geo_order.size(), &geo_order));
  EXPECT_EQ(geo_order.size(), 1U);
  EXPECT_EQ(geo_order[0], 0U);

  geo_order.push_back(0);
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply(",", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("2,", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("1", geo_order.size(), &geo_order));
  EXPECT_FALSE(
    download_mgr.ValidateGeoReply("3,2,1", geo_order.size(), &geo_order));
  EXPECT_TRUE(
    download_mgr.ValidateGeoReply("2,1", geo_order.size(), &geo_order));
  EXPECT_EQ(geo_order.size(), 2U);
  EXPECT_EQ(geo_order[0], 1U);
  EXPECT_EQ(geo_order[1], 0U);

  EXPECT_TRUE(
    download_mgr.ValidateGeoReply("2,1\n", geo_order.size(), &geo_order));
  EXPECT_EQ(geo_order.size(), 2U);
  EXPECT_EQ(geo_order[0], 1U);
  EXPECT_EQ(geo_order[1], 0U);

  geo_order.push_back(0);
  geo_order.push_back(0);
  EXPECT_TRUE(
    download_mgr.ValidateGeoReply("4,3,1,2\n", geo_order.size(), &geo_order));
  EXPECT_EQ(geo_order.size(), 4U);
  EXPECT_EQ(geo_order[0], 3U);
  EXPECT_EQ(geo_order[1], 2U);
  EXPECT_EQ(geo_order[2], 0U);
  EXPECT_EQ(geo_order[3], 1U);
}

}  // namespace download
<|endoftext|>