<commit_before>import lib.llvm.llvm;
import lib.llvm.llvm.ModuleRef;
import std._str;
import std._vec;
import std.os.target_os;
import util.common.istr;

const int wordsz = 4;

fn wstr(int i) -> str {
    ret istr(i * wordsz);
}

fn start() -> vec[str] {
    ret vec(".cfi_startproc");
}

fn end() -> vec[str] {
    ret vec(".cfi_endproc");
}

fn save_callee_saves() -> vec[str] {
    ret vec("pushl %ebp",
            "pushl %edi",
            "pushl %esi",
            "pushl %ebx");
}

fn save_callee_saves_with_cfi() -> vec[str] {
    auto offset = 8;
    auto t;
    t  = vec("pushl %ebp");
    t += vec(".cfi_def_cfa_offset " + istr(offset));
    t += vec(".cfi_offset 5, -" + istr(offset));

    t += vec("pushl %edi");
    offset += 4;
    t += vec(".cfi_def_cfa_offset " + istr(offset));

    t += vec("pushl %esi");
    offset += 4;
    t += vec(".cfi_def_cfa_offset " + istr(offset));

    t += vec("pushl %ebx");
    offset += 4;
    t += vec(".cfi_def_cfa_offset " + istr(offset));
    ret t;
}

fn restore_callee_saves() -> vec[str] {
    ret vec("popl  %ebx",
            "popl  %esi",
            "popl  %edi",
            "popl  %ebp");
}

fn load_esp_from_rust_sp_first_arg() -> vec[str] {
    ret vec("movl  " + wstr(abi.task_field_rust_sp) + "(%ecx), %esp");
}

fn load_esp_from_runtime_sp_first_arg() -> vec[str] {
    ret vec("movl  " + wstr(abi.task_field_runtime_sp) + "(%ecx), %esp");
}

fn store_esp_to_rust_sp_first_arg() -> vec[str] {
    ret vec("movl  %esp, " + wstr(abi.task_field_rust_sp) + "(%ecx)");
}

fn store_esp_to_runtime_sp_first_arg() -> vec[str] {
    ret vec("movl  %esp, " + wstr(abi.task_field_runtime_sp) + "(%ecx)");
}

fn load_esp_from_rust_sp_second_arg() -> vec[str] {
    ret vec("movl  " + wstr(abi.task_field_rust_sp) + "(%edx), %esp");
}

fn load_esp_from_runtime_sp_second_arg() -> vec[str] {
    ret vec("movl  " + wstr(abi.task_field_runtime_sp) + "(%edx), %esp");
}

fn store_esp_to_rust_sp_second_arg() -> vec[str] {
    ret vec("movl  %esp, " + wstr(abi.task_field_rust_sp) + "(%edx)");
}

fn store_esp_to_runtime_sp_second_arg() -> vec[str] {
    ret vec("movl  %esp, " + wstr(abi.task_field_runtime_sp) + "(%edx)");
}

/*
 * This is a bit of glue-code. It should be emitted once per
 * compilation unit.
 *
 *   - save regs on C stack
 *   - align sp on a 16-byte boundary
 *   - save sp to task.runtime_sp (runtime_sp is thus always aligned)
 *   - load saved task sp (switch stack)
 *   - restore saved task regs
 *   - return to saved task pc
 *
 * Our incoming stack looks like this:
 *
 *   *esp+4        = [arg1   ] = task ptr
 *   *esp          = [retpc  ]
 */

fn rust_activate_glue() -> vec[str] {
    ret vec("movl  4(%esp), %ecx    # ecx = rust_task")
        + save_callee_saves()
        + store_esp_to_runtime_sp_first_arg()
        + load_esp_from_rust_sp_first_arg()

        /*
         * There are two paths we can arrive at this code from:
         *
         *
         *   1. We are activating a task for the first time. When we switch
         *      into the task stack and 'ret' to its first instruction, we'll
         *      start doing whatever the first instruction says. Probably
         *      saving registers and starting to establish a frame. Harmless
         *      stuff, doesn't look at task->rust_sp again except when it
         *      clobbers it during a later native call.
         *
         *
         *   2. We are resuming a task that was descheduled by the yield glue
         *      below.  When we switch into the task stack and 'ret', we'll be
         *      ret'ing to a very particular instruction:
         *
         *              "esp <- task->rust_sp"
         *
         *      this is the first instruction we 'ret' to after this glue,
         *      because it is the first instruction following *any* native
         *      call, and the task we are activating was descheduled
         *      mid-native-call.
         *
         *      Unfortunately for us, we have already restored esp from
         *      task->rust_sp and are about to eat the 5 words off the top of
         *      it.
         *
         *
         *      | ...    | <-- where esp will be once we restore + ret, below,
         *      | retpc  |     and where we'd *like* task->rust_sp to wind up.
         *      | ebp    |
         *      | edi    |
         *      | esi    |
         *      | ebx    | <-- current task->rust_sp == current esp
         *
         *
         *      This is a problem. If we return to "esp <- task->rust_sp" it
         *      will push esp back down by 5 words. This manifests as a rust
         *      stack that grows by 5 words on each yield/reactivate. Not
         *      good.
         *
         *      So what we do here is just adjust task->rust_sp up 5 words as
         *      well, to mirror the movement in esp we're about to
         *      perform. That way the "esp <- task->rust_sp" we 'ret' to below
         *      will be a no-op. Esp won't move, and the task's stack won't
         *      grow.
         */
        + vec("addl  $20, " + wstr(abi.task_field_rust_sp) + "(%ecx)")

        /*
         * In most cases, the function we're returning to (activating)
         * will have saved any caller-saves before it yielded via native call,
         * so no work to do here. With one exception: when we're initially
         * activating, the task needs to be in the fastcall 2nd parameter
         * expected by the rust main function. That's edx.
         */
        + vec("mov  %ecx, %edx")

        + restore_callee_saves()
        + vec("ret");
}

/* More glue code, this time the 'bottom half' of yielding.
 *
 * We arrived here because an native call decided to deschedule the
 * running task. So the native call's return address got patched to the
 * first instruction of this glue code.
 *
 * When the native call does 'ret' it will come here, and its esp will be
 * pointing to the last argument pushed on the C stack before making
 * the native call: the 0th argument to the native call, which is always
 * the task ptr performing the native call. That's where we take over.
 *
 * Our goal is to complete the descheduling
 *
 *   - Switch over to the task stack temporarily.
 *
 *   - Save the task's callee-saves onto the task stack.
 *     (the task is now 'descheduled', safe to set aside)
 *
 *   - Switch *back* to the C stack.
 *
 *   - Restore the C-stack callee-saves.
 *
 *   - Return to the caller on the C stack that activated the task.
 *
 */

fn rust_yield_glue() -> vec[str] {
    ret vec("movl  0(%esp), %ecx    # ecx = rust_task")
        + load_esp_from_rust_sp_first_arg()
        + save_callee_saves()
        + store_esp_to_rust_sp_first_arg()
        + load_esp_from_runtime_sp_first_arg()
        + restore_callee_saves()
        + vec("ret");
}

fn native_glue(int n_args, abi.native_glue_type ngt) -> vec[str] {

    let bool pass_task;
    alt (ngt) {
        case (abi.ngt_rust)         { pass_task = true; }
        case (abi.ngt_pure_rust)    { pass_task = true; }
        case (abi.ngt_cdecl)        { pass_task = false; }
    }

    /*
     * 0, 4, 8, 12 are callee-saves
     * 16 is retpc
     * 20 .. (5+i) * 4 are args
     *
     * ecx is taskptr
     * edx is callee
     *
     */

    fn copy_arg(bool pass_task, uint i) -> str {
        if (i == 0u && pass_task) {
            ret "movl  %edx, (%esp)";
        }
        auto dst_off = wstr(0 + (i as int));
        auto src_off;
        if (pass_task) {
            src_off = wstr(4 + (i as int));
        } else {
            src_off = wstr(5 + (i as int));
        }
        auto m = vec("movl  " + src_off + "(%ebp),%eax",
                     "movl  %eax," + dst_off + "(%esp)");
        ret _str.connect(m, "\n\t");
    }

    auto carg = bind copy_arg(pass_task, _);

    ret
        start()
        + save_callee_saves_with_cfi()

        + vec("movl  %esp, %ebp     # ebp = rust_sp")
        + vec(".cfi_def_cfa_register 5")

        + store_esp_to_rust_sp_second_arg()
        + load_esp_from_runtime_sp_second_arg()

        + vec("subl  $" + wstr(n_args) + ", %esp   # esp -= args",
              "andl  $~0xf, %esp    # align esp down")

        + _vec.init_fn[str](carg, (n_args) as uint)

        +  vec("movl  %edx, %edi     # save task from edx to edi",
               "call  *%ecx          # call *%ecx",
               "movl  %edi, %edx     # restore edi-saved task to edx")

        + load_esp_from_rust_sp_second_arg()
        + restore_callee_saves()
        + vec("ret")
        + end();

}

fn decl_glue(int align, str prefix, str name, vec[str] insns) -> str {
    auto sym = prefix + name;
    ret "\t.globl " + sym + "\n" +
        "\t.balign " + istr(align) + "\n" +
        sym + ":\n" +
        "\t" + _str.connect(insns, "\n\t");
}

fn decl_native_glue(int align, str prefix, abi.native_glue_type ngt, uint n)
        -> str {
    let int i = n as int;
    ret decl_glue(align, prefix,
                  abi.native_glue_name(i, ngt),
                  native_glue(i, ngt));
}

fn get_symbol_prefix() -> str {
    if (_str.eq(target_os(), "macos") ||
        _str.eq(target_os(), "win32")) {
        ret "_";
    } else {
        ret "";
    }
}

fn get_module_asm() -> str {
    auto align = 4;

    auto prefix = get_symbol_prefix();

    auto glues =
        vec(decl_glue(align, prefix,
                      abi.activate_glue_name(),
                      rust_activate_glue()),

            decl_glue(align, prefix,
                      abi.yield_glue_name(),
                      rust_yield_glue()))

        + _vec.init_fn[str](bind decl_native_glue(align, prefix,
            abi.ngt_rust, _), (abi.n_native_glues + 1) as uint)
        + _vec.init_fn[str](bind decl_native_glue(align, prefix,
            abi.ngt_pure_rust, _), (abi.n_native_glues + 1) as uint)
        + _vec.init_fn[str](bind decl_native_glue(align, prefix,
            abi.ngt_cdecl, _), (abi.n_native_glues + 1) as uint);

    ret _str.connect(glues, "\n\n");
}

fn get_meta_sect_name() -> str {
    if (_str.eq(target_os(), "macos")) {
        ret "__DATA,__note.rustc";
    }
    if (_str.eq(target_os(), "win32")) {
        ret ".note.rustc";
    }
    ret ".note.rustc";
}

fn get_data_layout() -> str {
    if (_str.eq(target_os(), "macos")) {
      ret "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64" +
        "-f32:32:32-f64:32:64-v64:64:64-v128:128:128-a0:0:64-f80:128:128" +
        "-n8:16:32";
    }
    if (_str.eq(target_os(), "win32")) {
      ret "e-p:32:32-f64:64:64-i64:64:64-f80:32:32-n8:16:32";
    }
    ret "e-p:32:32-f64:32:64-i64:32:64-f80:32:32-n8:16:32";
}

fn get_target_triple() -> str {
    if (_str.eq(target_os(), "macos")) {
        ret "i686-apple-darwin";
    }
    if (_str.eq(target_os(), "win32")) {
        ret "i686-pc-mingw32";
    }
    ret "i686-pc-linux-gnu";
}

//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C $RBUILD 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
//
<commit_msg>Use symbolic register names so that we get the correct encoding on OS X.<commit_after>import lib.llvm.llvm;
import lib.llvm.llvm.ModuleRef;
import std._str;
import std._vec;
import std.os.target_os;
import util.common.istr;

const int wordsz = 4;

fn wstr(int i) -> str {
    ret istr(i * wordsz);
}

fn start() -> vec[str] {
    ret vec(".cfi_startproc");
}

fn end() -> vec[str] {
    ret vec(".cfi_endproc");
}

fn save_callee_saves() -> vec[str] {
    ret vec("pushl %ebp",
            "pushl %edi",
            "pushl %esi",
            "pushl %ebx");
}

fn save_callee_saves_with_cfi() -> vec[str] {
    auto offset = 8;
    auto t;
    t  = vec("pushl %ebp");
    t += vec(".cfi_def_cfa_offset " + istr(offset));
    t += vec(".cfi_offset %ebp, -" + istr(offset));

    t += vec("pushl %edi");
    offset += 4;
    t += vec(".cfi_def_cfa_offset " + istr(offset));

    t += vec("pushl %esi");
    offset += 4;
    t += vec(".cfi_def_cfa_offset " + istr(offset));

    t += vec("pushl %ebx");
    offset += 4;
    t += vec(".cfi_def_cfa_offset " + istr(offset));
    ret t;
}

fn restore_callee_saves() -> vec[str] {
    ret vec("popl  %ebx",
            "popl  %esi",
            "popl  %edi",
            "popl  %ebp");
}

fn load_esp_from_rust_sp_first_arg() -> vec[str] {
    ret vec("movl  " + wstr(abi.task_field_rust_sp) + "(%ecx), %esp");
}

fn load_esp_from_runtime_sp_first_arg() -> vec[str] {
    ret vec("movl  " + wstr(abi.task_field_runtime_sp) + "(%ecx), %esp");
}

fn store_esp_to_rust_sp_first_arg() -> vec[str] {
    ret vec("movl  %esp, " + wstr(abi.task_field_rust_sp) + "(%ecx)");
}

fn store_esp_to_runtime_sp_first_arg() -> vec[str] {
    ret vec("movl  %esp, " + wstr(abi.task_field_runtime_sp) + "(%ecx)");
}

fn load_esp_from_rust_sp_second_arg() -> vec[str] {
    ret vec("movl  " + wstr(abi.task_field_rust_sp) + "(%edx), %esp");
}

fn load_esp_from_runtime_sp_second_arg() -> vec[str] {
    ret vec("movl  " + wstr(abi.task_field_runtime_sp) + "(%edx), %esp");
}

fn store_esp_to_rust_sp_second_arg() -> vec[str] {
    ret vec("movl  %esp, " + wstr(abi.task_field_rust_sp) + "(%edx)");
}

fn store_esp_to_runtime_sp_second_arg() -> vec[str] {
    ret vec("movl  %esp, " + wstr(abi.task_field_runtime_sp) + "(%edx)");
}

/*
 * This is a bit of glue-code. It should be emitted once per
 * compilation unit.
 *
 *   - save regs on C stack
 *   - align sp on a 16-byte boundary
 *   - save sp to task.runtime_sp (runtime_sp is thus always aligned)
 *   - load saved task sp (switch stack)
 *   - restore saved task regs
 *   - return to saved task pc
 *
 * Our incoming stack looks like this:
 *
 *   *esp+4        = [arg1   ] = task ptr
 *   *esp          = [retpc  ]
 */

fn rust_activate_glue() -> vec[str] {
    ret vec("movl  4(%esp), %ecx    # ecx = rust_task")
        + save_callee_saves()
        + store_esp_to_runtime_sp_first_arg()
        + load_esp_from_rust_sp_first_arg()

        /*
         * There are two paths we can arrive at this code from:
         *
         *
         *   1. We are activating a task for the first time. When we switch
         *      into the task stack and 'ret' to its first instruction, we'll
         *      start doing whatever the first instruction says. Probably
         *      saving registers and starting to establish a frame. Harmless
         *      stuff, doesn't look at task->rust_sp again except when it
         *      clobbers it during a later native call.
         *
         *
         *   2. We are resuming a task that was descheduled by the yield glue
         *      below.  When we switch into the task stack and 'ret', we'll be
         *      ret'ing to a very particular instruction:
         *
         *              "esp <- task->rust_sp"
         *
         *      this is the first instruction we 'ret' to after this glue,
         *      because it is the first instruction following *any* native
         *      call, and the task we are activating was descheduled
         *      mid-native-call.
         *
         *      Unfortunately for us, we have already restored esp from
         *      task->rust_sp and are about to eat the 5 words off the top of
         *      it.
         *
         *
         *      | ...    | <-- where esp will be once we restore + ret, below,
         *      | retpc  |     and where we'd *like* task->rust_sp to wind up.
         *      | ebp    |
         *      | edi    |
         *      | esi    |
         *      | ebx    | <-- current task->rust_sp == current esp
         *
         *
         *      This is a problem. If we return to "esp <- task->rust_sp" it
         *      will push esp back down by 5 words. This manifests as a rust
         *      stack that grows by 5 words on each yield/reactivate. Not
         *      good.
         *
         *      So what we do here is just adjust task->rust_sp up 5 words as
         *      well, to mirror the movement in esp we're about to
         *      perform. That way the "esp <- task->rust_sp" we 'ret' to below
         *      will be a no-op. Esp won't move, and the task's stack won't
         *      grow.
         */
        + vec("addl  $20, " + wstr(abi.task_field_rust_sp) + "(%ecx)")

        /*
         * In most cases, the function we're returning to (activating)
         * will have saved any caller-saves before it yielded via native call,
         * so no work to do here. With one exception: when we're initially
         * activating, the task needs to be in the fastcall 2nd parameter
         * expected by the rust main function. That's edx.
         */
        + vec("mov  %ecx, %edx")

        + restore_callee_saves()
        + vec("ret");
}

/* More glue code, this time the 'bottom half' of yielding.
 *
 * We arrived here because an native call decided to deschedule the
 * running task. So the native call's return address got patched to the
 * first instruction of this glue code.
 *
 * When the native call does 'ret' it will come here, and its esp will be
 * pointing to the last argument pushed on the C stack before making
 * the native call: the 0th argument to the native call, which is always
 * the task ptr performing the native call. That's where we take over.
 *
 * Our goal is to complete the descheduling
 *
 *   - Switch over to the task stack temporarily.
 *
 *   - Save the task's callee-saves onto the task stack.
 *     (the task is now 'descheduled', safe to set aside)
 *
 *   - Switch *back* to the C stack.
 *
 *   - Restore the C-stack callee-saves.
 *
 *   - Return to the caller on the C stack that activated the task.
 *
 */

fn rust_yield_glue() -> vec[str] {
    ret vec("movl  0(%esp), %ecx    # ecx = rust_task")
        + load_esp_from_rust_sp_first_arg()
        + save_callee_saves()
        + store_esp_to_rust_sp_first_arg()
        + load_esp_from_runtime_sp_first_arg()
        + restore_callee_saves()
        + vec("ret");
}

fn native_glue(int n_args, abi.native_glue_type ngt) -> vec[str] {

    let bool pass_task;
    alt (ngt) {
        case (abi.ngt_rust)         { pass_task = true; }
        case (abi.ngt_pure_rust)    { pass_task = true; }
        case (abi.ngt_cdecl)        { pass_task = false; }
    }

    /*
     * 0, 4, 8, 12 are callee-saves
     * 16 is retpc
     * 20 .. (5+i) * 4 are args
     *
     * ecx is taskptr
     * edx is callee
     *
     */

    fn copy_arg(bool pass_task, uint i) -> str {
        if (i == 0u && pass_task) {
            ret "movl  %edx, (%esp)";
        }
        auto dst_off = wstr(0 + (i as int));
        auto src_off;
        if (pass_task) {
            src_off = wstr(4 + (i as int));
        } else {
            src_off = wstr(5 + (i as int));
        }
        auto m = vec("movl  " + src_off + "(%ebp),%eax",
                     "movl  %eax," + dst_off + "(%esp)");
        ret _str.connect(m, "\n\t");
    }

    auto carg = bind copy_arg(pass_task, _);

    ret
        start()
        + save_callee_saves_with_cfi()

        + vec("movl  %esp, %ebp     # ebp = rust_sp")
        + vec(".cfi_def_cfa_register %ebp")

        + store_esp_to_rust_sp_second_arg()
        + load_esp_from_runtime_sp_second_arg()

        + vec("subl  $" + wstr(n_args) + ", %esp   # esp -= args",
              "andl  $~0xf, %esp    # align esp down")

        + _vec.init_fn[str](carg, (n_args) as uint)

        +  vec("movl  %edx, %edi     # save task from edx to edi",
               "call  *%ecx          # call *%ecx",
               "movl  %edi, %edx     # restore edi-saved task to edx")

        + load_esp_from_rust_sp_second_arg()
        + restore_callee_saves()
        + vec("ret")
        + end();

}

fn decl_glue(int align, str prefix, str name, vec[str] insns) -> str {
    auto sym = prefix + name;
    ret "\t.globl " + sym + "\n" +
        "\t.balign " + istr(align) + "\n" +
        sym + ":\n" +
        "\t" + _str.connect(insns, "\n\t");
}

fn decl_native_glue(int align, str prefix, abi.native_glue_type ngt, uint n)
        -> str {
    let int i = n as int;
    ret decl_glue(align, prefix,
                  abi.native_glue_name(i, ngt),
                  native_glue(i, ngt));
}

fn get_symbol_prefix() -> str {
    if (_str.eq(target_os(), "macos") ||
        _str.eq(target_os(), "win32")) {
        ret "_";
    } else {
        ret "";
    }
}

fn get_module_asm() -> str {
    auto align = 4;

    auto prefix = get_symbol_prefix();

    auto glues =
        vec(decl_glue(align, prefix,
                      abi.activate_glue_name(),
                      rust_activate_glue()),

            decl_glue(align, prefix,
                      abi.yield_glue_name(),
                      rust_yield_glue()))

        + _vec.init_fn[str](bind decl_native_glue(align, prefix,
            abi.ngt_rust, _), (abi.n_native_glues + 1) as uint)
        + _vec.init_fn[str](bind decl_native_glue(align, prefix,
            abi.ngt_pure_rust, _), (abi.n_native_glues + 1) as uint)
        + _vec.init_fn[str](bind decl_native_glue(align, prefix,
            abi.ngt_cdecl, _), (abi.n_native_glues + 1) as uint);

    ret _str.connect(glues, "\n\n");
}

fn get_meta_sect_name() -> str {
    if (_str.eq(target_os(), "macos")) {
        ret "__DATA,__note.rustc";
    }
    if (_str.eq(target_os(), "win32")) {
        ret ".note.rustc";
    }
    ret ".note.rustc";
}

fn get_data_layout() -> str {
    if (_str.eq(target_os(), "macos")) {
      ret "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64" +
        "-f32:32:32-f64:32:64-v64:64:64-v128:128:128-a0:0:64-f80:128:128" +
        "-n8:16:32";
    }
    if (_str.eq(target_os(), "win32")) {
      ret "e-p:32:32-f64:64:64-i64:64:64-f80:32:32-n8:16:32";
    }
    ret "e-p:32:32-f64:32:64-i64:32:64-f80:32:32-n8:16:32";
}

fn get_target_triple() -> str {
    if (_str.eq(target_os(), "macos")) {
        ret "i686-apple-darwin";
    }
    if (_str.eq(target_os(), "win32")) {
        ret "i686-pc-mingw32";
    }
    ret "i686-pc-linux-gnu";
}

//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C $RBUILD 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
//
<|endoftext|>


<commit_before>// Copyright 2016 Mozilla Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![allow(unused_imports, dead_code, unused_variables)]

use crate::compiler::args::*;
use crate::compiler::c::{CCompilerImpl, CCompilerKind, Language, ParsedArguments};
use crate::compiler::gcc::ArgData::*;
use crate::compiler::{gcc, write_temp_file, Cacheable, CompileCommand, CompilerArguments};
use crate::dist;
use crate::mock_command::{CommandCreator, CommandCreatorSync, RunCommand};
use crate::util::{run_input_output, OsStrExt};
use futures::future::{self, Future};
use futures_cpupool::CpuPool;
use log::Level::Trace;
use std::ffi::OsString;
use std::fs::File;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process;

use crate::errors::*;

/// A unit struct on which to implement `CCompilerImpl`.
#[derive(Clone, Debug)]
pub struct NVCC;

impl CCompilerImpl for NVCC {
    fn kind(&self) -> CCompilerKind {
        CCompilerKind::NVCC
    }
    fn parse_arguments(
        &self,
        arguments: &[OsString],
        cwd: &Path,
    ) -> CompilerArguments<ParsedArguments> {
        gcc::parse_arguments(arguments, cwd, (&gcc::ARGS[..], &ARGS[..]))
    }

    fn preprocess<T>(
        &self,
        creator: &T,
        executable: &Path,
        parsed_args: &ParsedArguments,
        cwd: &Path,
        env_vars: &[(OsString, OsString)],
        may_dist: bool,
        rewrite_includes_only: bool,
    ) -> SFuture<process::Output>
    where
        T: CommandCreatorSync,
    {
        trace!("preprocess");
        let mut cmd = creator.clone().new_command_sync(executable);

        //NVCC only supports `-E` when it comes after preprocessor
        //and common flags
        cmd.args(&parsed_args.preprocessor_args);
        cmd.args(&parsed_args.common_args);

        //We need to add "-rdc=true" if we are compiling with `-dc`
        //So that the preprocessor has the correct implicit defines
        if parsed_args.compilation_flag == "-dc" {
            cmd.arg("-rdc=true");
        }

        let language = match parsed_args.language {
            Language::C => "c",
            Language::Cxx => "c++",
            Language::ObjectiveC => "objective-c",
            Language::ObjectiveCxx => "objective-c++",
            Language::Cuda => "cu",
        };

        cmd.arg("-x").arg(language)
           .arg(&parsed_args.input)
           .arg("-E")
           .env_clear()
           .envs(env_vars.iter().map(|&(ref k, ref v)| (k, v)))
           .current_dir(cwd);
        if log_enabled!(Trace) {
            trace!("preprocess: {:?}", cmd);
        }
        run_input_output(cmd, None)
    }

    fn generate_compile_commands(
        &self,
        path_transformer: &mut dist::PathTransformer,
        executable: &Path,
        parsed_args: &ParsedArguments,
        cwd: &Path,
        env_vars: &[(OsString, OsString)],
        rewrite_includes_only: bool,
    ) -> Result<(CompileCommand, Option<dist::CompileCommand>, Cacheable)> {
        gcc::generate_compile_commands(
            path_transformer,
            executable,
            parsed_args,
            cwd,
            env_vars,
            self.kind(),
            rewrite_includes_only,
        )
    }
}

counted_array!(pub static ARGS: [ArgInfo<gcc::ArgData>; _] = [
    // take_arg!("--compiler-bindir", PathBuf, Separated, ExtraHashFile),
    // take_arg!("-ccbin", PathBuf, Separated, ExtraHashFile),

    take_arg!("--archive-options options", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--compiler-options", OsString, CanBeSeparated('='), PreprocessorArgument),
    take_arg!("--generate-code", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--gpu-architecture", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--gpu-code", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--linker-options", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--maxrregcount", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--nvlink-options", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--ptxas-options", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--relocatable-device-code", OsString, CanBeSeparated('='), PreprocessorArgument),

    take_arg!("-Xarchive", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-Xcompiler", OsString, CanBeSeparated('='), PreprocessorArgument),
    take_arg!("-Xlinker", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-Xnvlink", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-Xptxas", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-arch", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-code", OsString, CanBeSeparated('='), PassThrough),
    flag!("-dc", DoCompilation),
    take_arg!("-gencode", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-maxrregcount", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-rdc", OsString, CanBeSeparated('='), PreprocessorArgument),
    take_arg!("-x", OsString, CanBeSeparated('='), Language),

]);

#[cfg(test)]
mod test {
    use super::*;
    use crate::compiler::gcc;
    use crate::compiler::*;
    use crate::mock_command::*;
    use crate::test::utils::*;
    use futures::Future;
    use futures_cpupool::CpuPool;
    use std::collections::HashMap;
    use std::path::PathBuf;

    fn parse_arguments_(arguments: Vec<String>) -> CompilerArguments<ParsedArguments> {
        let arguments = arguments.iter().map(OsString::from).collect::<Vec<_>>();
        NVCC.parse_arguments(&arguments, ".".as_ref())
    }

    macro_rules! parses {
        ( $( $s:expr ),* ) => {
            match parse_arguments_(vec![ $( $s.to_string(), )* ]) {
                CompilerArguments::Ok(a) => a,
                o => panic!("Got unexpected parse result: {:?}", o),
            }
        }
    }

    #[test]
    fn test_parse_arguments_simple_c() {
        let a = parses!("-c", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::C, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert!(a.common_args.is_empty());
    }

    #[test]
    fn test_parse_arguments_simple_cu() {
        let a = parses!("-c", "foo.cu", "-o", "foo.o");
        assert_eq!(Some("foo.cu"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert!(a.common_args.is_empty());
    }

    #[test]
    fn test_parse_arguments_simple_c_as_cu() {
        let a = parses!("-x","cu", "-c", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert!(a.common_args.is_empty());
    }

    #[test]
    fn test_parse_arguments_dc_compile_flag() {
        let a = parses!("-x","cu", "-dc", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_eq!(Some("-dc"), a.compilation_flag.to_str());
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert!(a.common_args.is_empty());
    }

    #[test]
    fn test_parse_md_mt_flags_cu() {
        let a = parses!("-x", "cu", "-c", "foo.c", "-fabc",
                           "-MD", "-MT", "foo.o", "-MF", "foo.o.d", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_eq!(Some("-c"), a.compilation_flag.to_str());
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert_eq!(
            ovec!["-MD", "-MF", "foo.o.d", "-MT", "foo.o"],
            a.preprocessor_args
        );
        assert_eq!(ovec!["-fabc"], a.common_args);
    }

    #[test]
    fn test_parse_generate_code_flags() {
        let a = parses!("-x","cu", "--generate-code=arch=compute_61,code=sm_61",
                        "-c", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert_eq!(
            ovec!["--generate-code", "arch=compute_61,code=sm_61"],
            a.common_args
        );
    }

    #[test]
    fn test_parse_pass_to_host_flags() {
        let a = parses!("-x=cu",
                        "--generate-code=arch=compute_60,code=[sm_60,sm_61]",
                        "-Xnvlink=--suppress-stack-size-warning",
                        "-Xcompiler", "-fPIC,-fno-common",
                        "-Xcompiler=-fvisibility=hidden",
                        "-Xcompiler=-Wall,-Wno-unknown-pragmas,-Wno-unused-local-typedefs",
                        "-Xcudafe", "--display_error_number",
                        "-c", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert_eq!(
            ovec!["-Xcompiler", "-fPIC,-fno-common", "-Xcompiler", "-fvisibility=hidden",
                  "-Xcompiler", "-Wall,-Wno-unknown-pragmas,-Wno-unused-local-typedefs"],
            a.preprocessor_args
        );
        assert_eq!(
            ovec!["--generate-code", "arch=compute_60,code=[sm_60,sm_61]",
                  "-Xnvlink", "--suppress-stack-size-warning", "-Xcudafe",
                  "--display_error_number"],
            a.common_args
        );
    }

}
<commit_msg>nvcc handles '=' with --include-path and -isystem<commit_after>// Copyright 2016 Mozilla Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![allow(unused_imports, dead_code, unused_variables)]

use crate::compiler::args::*;
use crate::compiler::c::{CCompilerImpl, CCompilerKind, Language, ParsedArguments};
use crate::compiler::gcc::ArgData::*;
use crate::compiler::{gcc, write_temp_file, Cacheable, CompileCommand, CompilerArguments};
use crate::dist;
use crate::mock_command::{CommandCreator, CommandCreatorSync, RunCommand};
use crate::util::{run_input_output, OsStrExt};
use futures::future::{self, Future};
use futures_cpupool::CpuPool;
use log::Level::Trace;
use std::ffi::OsString;
use std::fs::File;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process;

use crate::errors::*;

/// A unit struct on which to implement `CCompilerImpl`.
#[derive(Clone, Debug)]
pub struct NVCC;

impl CCompilerImpl for NVCC {
    fn kind(&self) -> CCompilerKind {
        CCompilerKind::NVCC
    }
    fn parse_arguments(
        &self,
        arguments: &[OsString],
        cwd: &Path,
    ) -> CompilerArguments<ParsedArguments> {
        gcc::parse_arguments(arguments, cwd, (&gcc::ARGS[..], &ARGS[..]))
    }

    fn preprocess<T>(
        &self,
        creator: &T,
        executable: &Path,
        parsed_args: &ParsedArguments,
        cwd: &Path,
        env_vars: &[(OsString, OsString)],
        may_dist: bool,
        rewrite_includes_only: bool,
    ) -> SFuture<process::Output>
    where
        T: CommandCreatorSync,
    {
        trace!("preprocess");
        let mut cmd = creator.clone().new_command_sync(executable);

        //NVCC only supports `-E` when it comes after preprocessor
        //and common flags
        cmd.args(&parsed_args.preprocessor_args);
        cmd.args(&parsed_args.common_args);

        //We need to add "-rdc=true" if we are compiling with `-dc`
        //So that the preprocessor has the correct implicit defines
        if parsed_args.compilation_flag == "-dc" {
            cmd.arg("-rdc=true");
        }

        let language = match parsed_args.language {
            Language::C => "c",
            Language::Cxx => "c++",
            Language::ObjectiveC => "objective-c",
            Language::ObjectiveCxx => "objective-c++",
            Language::Cuda => "cu",
        };

        cmd.arg("-x").arg(language)
           .arg(&parsed_args.input)
           .arg("-E")
           .env_clear()
           .envs(env_vars.iter().map(|&(ref k, ref v)| (k, v)))
           .current_dir(cwd);
        if log_enabled!(Trace) {
            trace!("preprocess: {:?}", cmd);
        }
        run_input_output(cmd, None)
    }

    fn generate_compile_commands(
        &self,
        path_transformer: &mut dist::PathTransformer,
        executable: &Path,
        parsed_args: &ParsedArguments,
        cwd: &Path,
        env_vars: &[(OsString, OsString)],
        rewrite_includes_only: bool,
    ) -> Result<(CompileCommand, Option<dist::CompileCommand>, Cacheable)> {
        gcc::generate_compile_commands(
            path_transformer,
            executable,
            parsed_args,
            cwd,
            env_vars,
            self.kind(),
            rewrite_includes_only,
        )
    }
}

counted_array!(pub static ARGS: [ArgInfo<gcc::ArgData>; _] = [
    // take_arg!("--compiler-bindir", PathBuf, Separated, ExtraHashFile),
    // take_arg!("-ccbin", PathBuf, Separated, ExtraHashFile),

    take_arg!("--archive-options options", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--compiler-options", OsString, CanBeSeparated('='), PreprocessorArgument),
    take_arg!("--generate-code", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--gpu-architecture", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--gpu-code", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--include-path", PathBuf, CanBeSeparated('='), PreprocessorArgumentPath),
    take_arg!("--linker-options", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--maxrregcount", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--nvlink-options", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--ptxas-options", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("--relocatable-device-code", OsString, CanBeSeparated('='), PreprocessorArgument),
    take_arg!("--system-include", PathBuf, CanBeSeparated('='), PreprocessorArgumentPath),

    take_arg!("-Xarchive", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-Xcompiler", OsString, CanBeSeparated('='), PreprocessorArgument),
    take_arg!("-Xlinker", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-Xnvlink", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-Xptxas", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-arch", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-code", OsString, CanBeSeparated('='), PassThrough),
    flag!("-dc", DoCompilation),
    take_arg!("-gencode", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-isystem", PathBuf, CanBeSeparated('='), PreprocessorArgumentPath),
    take_arg!("-maxrregcount", OsString, CanBeSeparated('='), PassThrough),
    take_arg!("-rdc", OsString, CanBeSeparated('='), PreprocessorArgument),
    take_arg!("-x", OsString, CanBeSeparated('='), Language),

]);

#[cfg(test)]
mod test {
    use super::*;
    use crate::compiler::gcc;
    use crate::compiler::*;
    use crate::mock_command::*;
    use crate::test::utils::*;
    use futures::Future;
    use futures_cpupool::CpuPool;
    use std::collections::HashMap;
    use std::path::PathBuf;

    fn parse_arguments_(arguments: Vec<String>) -> CompilerArguments<ParsedArguments> {
        let arguments = arguments.iter().map(OsString::from).collect::<Vec<_>>();
        NVCC.parse_arguments(&arguments, ".".as_ref())
    }

    macro_rules! parses {
        ( $( $s:expr ),* ) => {
            match parse_arguments_(vec![ $( $s.to_string(), )* ]) {
                CompilerArguments::Ok(a) => a,
                o => panic!("Got unexpected parse result: {:?}", o),
            }
        }
    }

    #[test]
    fn test_parse_arguments_simple_c() {
        let a = parses!("-c", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::C, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert!(a.common_args.is_empty());
    }

    #[test]
    fn test_parse_arguments_simple_cu() {
        let a = parses!("-c", "foo.cu", "-o", "foo.o");
        assert_eq!(Some("foo.cu"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert!(a.common_args.is_empty());
    }

    #[test]
    fn test_parse_arguments_simple_c_as_cu() {
        let a = parses!("-x","cu", "-c", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert!(a.common_args.is_empty());
    }

    #[test]
    fn test_parse_arguments_dc_compile_flag() {
        let a = parses!("-x","cu", "-dc", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_eq!(Some("-dc"), a.compilation_flag.to_str());
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert!(a.common_args.is_empty());
    }

    #[test]
    fn test_parse_arguments_values() {
        let a = parses!("-c", "foo.cpp", "-fabc",
                        "-I", "include-file", "-o", "foo.o",
                        "--include-path", "include-file",
                        "-isystem=/system/include/file");
        assert_eq!(Some("foo.cpp"), a.input.to_str());
        assert_eq!(Language::Cxx, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert_eq!(
            ovec!["-Iinclude-file", "--include-path", "include-file", "-isystem", "/system/include/file"],
            a.preprocessor_args
        );
        assert!(a.dependency_args.is_empty());
        assert_eq!(ovec!["-fabc"], a.common_args);
    }

    #[test]
    fn test_parse_md_mt_flags_cu() {
        let a = parses!("-x", "cu", "-c", "foo.c", "-fabc",
                           "-MD", "-MT", "foo.o", "-MF", "foo.o.d", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_eq!(Some("-c"), a.compilation_flag.to_str());
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert_eq!(
            ovec!["-MD", "-MF", "foo.o.d", "-MT", "foo.o"],
            a.preprocessor_args
        );
        assert_eq!(ovec!["-fabc"], a.common_args);
    }

    #[test]
    fn test_parse_generate_code_flags() {
        let a = parses!("-x","cu", "--generate-code=arch=compute_61,code=sm_61",
                        "-c", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert!(a.preprocessor_args.is_empty());
        assert_eq!(
            ovec!["--generate-code", "arch=compute_61,code=sm_61"],
            a.common_args
        );
    }

    #[test]
    fn test_parse_pass_to_host_flags() {
        let a = parses!("-x=cu",
                        "--generate-code=arch=compute_60,code=[sm_60,sm_61]",
                        "-Xnvlink=--suppress-stack-size-warning",
                        "-Xcompiler", "-fPIC,-fno-common",
                        "-Xcompiler=-fvisibility=hidden",
                        "-Xcompiler=-Wall,-Wno-unknown-pragmas,-Wno-unused-local-typedefs",
                        "-Xcudafe", "--display_error_number",
                        "-c", "foo.c", "-o", "foo.o");
        assert_eq!(Some("foo.c"), a.input.to_str());
        assert_eq!(Language::Cuda, a.language);
        assert_map_contains!(a.outputs, ("obj", PathBuf::from("foo.o")));
        //TODO: fix assert_map_contains to assert no extra keys!
        assert_eq!(1, a.outputs.len());
        assert_eq!(
            ovec!["-Xcompiler", "-fPIC,-fno-common", "-Xcompiler", "-fvisibility=hidden",
                  "-Xcompiler", "-Wall,-Wno-unknown-pragmas,-Wno-unused-local-typedefs"],
            a.preprocessor_args
        );
        assert_eq!(
            ovec!["--generate-code", "arch=compute_60,code=[sm_60,sm_61]",
                  "-Xnvlink", "--suppress-stack-size-warning", "-Xcudafe",
                  "--display_error_number"],
            a.common_args
        );
    }

}
<|endoftext|>


<commit_before>use std::mem::size_of;
use std::os::raw::{c_int, c_long, c_longlong, c_short, c_uint, c_ulong, c_ulonglong, c_ushort};

use num_traits::{Bounded, Zero};
use pyo3::{ffi, prelude::*, pyobject_native_type_core, types::PyType, AsPyPointer, PyNativeType};

use crate::npyffi::{NpyTypes, PyArray_Descr, NPY_TYPES, PY_ARRAY_API};

pub use num_complex::{Complex32, Complex64};

/// Binding of [`numpy.dtype`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.html).
///
/// # Example
/// ```
/// use pyo3::types::IntoPyDict;
/// pyo3::Python::with_gil(|py| {
///     let locals = [("np", numpy::get_array_module(py).unwrap())].into_py_dict(py);
///     let dtype: &numpy::PyArrayDescr = py
///         .eval("np.array([1, 2, 3.0]).dtype", Some(locals), None)
///         .unwrap()
///         .downcast()
///         .unwrap();
///     assert!(dtype.is_equiv_to(numpy::dtype::<f64>(py)));
/// });
/// ```
pub struct PyArrayDescr(PyAny);

pyobject_native_type_core!(
    PyArrayDescr,
    *PY_ARRAY_API.get_type_object(NpyTypes::PyArrayDescr_Type),
    #module=Some("numpy"),
    #checkfunction=arraydescr_check
);

unsafe fn arraydescr_check(op: *mut ffi::PyObject) -> c_int {
    ffi::PyObject_TypeCheck(
        op,
        PY_ARRAY_API.get_type_object(NpyTypes::PyArrayDescr_Type),
    )
}

/// Returns the type descriptor ("dtype") for a registered type.
pub fn dtype<T: Element>(py: Python) -> &PyArrayDescr {
    T::get_dtype(py)
}

impl PyArrayDescr {
    /// Returns `self` as `*mut PyArray_Descr`.
    pub fn as_dtype_ptr(&self) -> *mut PyArray_Descr {
        self.as_ptr() as _
    }

    /// Returns `self` as `*mut PyArray_Descr` while increasing the reference count.
    ///
    /// Useful in cases where the descriptor is stolen by the API.
    pub fn into_dtype_ptr(&self) -> *mut PyArray_Descr {
        self.into_ptr() as _
    }

    /// Returns the internal `PyType` that this `dtype` holds.
    ///
    /// # Example
    /// ```
    /// pyo3::Python::with_gil(|py| {
    ///    let array = numpy::PyArray::from_vec(py, vec![0.0, 1.0, 2.0f64]);
    ///    let dtype = array.dtype();
    ///    assert_eq!(dtype.get_type().name().unwrap().to_string(), "float64");
    /// });
    /// ```
    pub fn get_type(&self) -> &PyType {
        let dtype_type_ptr = unsafe { *self.as_dtype_ptr() }.typeobj;
        unsafe { PyType::from_type_ptr(self.py(), dtype_type_ptr) }
    }

    /// Shortcut for creating a descriptor of 'object' type.
    pub fn object(py: Python) -> &Self {
        Self::from_npy_type(py, NPY_TYPES::NPY_OBJECT)
    }

    /// Returns the type descriptor ("dtype") for a registered type.
    pub fn of<T: Element>(py: Python) -> &Self {
        T::get_dtype(py)
    }

    /// Returns true if two type descriptors are equivalent.
    pub fn is_equiv_to(&self, other: &Self) -> bool {
        unsafe { PY_ARRAY_API.PyArray_EquivTypes(self.as_dtype_ptr(), other.as_dtype_ptr()) != 0 }
    }

    fn from_npy_type(py: Python, npy_type: NPY_TYPES) -> &Self {
        unsafe {
            let descr = PY_ARRAY_API.PyArray_DescrFromType(npy_type as _);
            py.from_owned_ptr(descr as _)
        }
    }

    /// Retrieves the
    /// [enumerated type](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types)
    /// for this type descriptor.
    pub fn get_typenum(&self) -> c_int {
        unsafe { *self.as_dtype_ptr() }.type_num
    }
}

/// Represents that a type can be an element of `PyArray`.
///
/// Currently, only integer/float/complex types are supported.
/// If you come up with a nice implementation for some other types, we're happy to receive your PR :)
/// You may refer to the [numpy document](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types)
/// for all types that numpy supports.
///
/// # Safety
///
/// A type `T` that implements this trait should be safe when managed in numpy
/// array, thus implementing this trait is marked unsafe. Data types that don't
/// contain Python objects (i.e., either the object type itself or record types
/// containing object-type fields) are assumed to be trivially copyable, which
/// is reflected in the `IS_COPY` flag. Furthermore, it is assumed that for
/// the object type the elements are pointers into the Python heap and that the
/// corresponding `Clone` implemenation will never panic as it only increases
/// the reference count.
///
/// # Custom element types
///
/// You can implement this trait to manage arrays of custom element types, but they still need to be stored
/// on Python's heap using PyO3's [Py](pyo3::Py) type.
///
/// ```
/// use numpy::{ndarray::Array2, Element, PyArray, PyArrayDescr, ToPyArray};
/// use pyo3::{pyclass, Py, Python};
///
/// #[pyclass]
/// pub struct CustomElement;
///
/// // The transparent wrapper is necessary as one cannot implement
/// // a foreign trait (`Element`) on a foreign type (`Py`) directly.
/// #[derive(Clone)]
/// #[repr(transparent)]
/// pub struct Wrapper(pub Py<CustomElement>);
///
/// unsafe impl Element for Wrapper {
///     const IS_COPY: bool = false;
///
///     fn get_dtype(py: Python) -> &PyArrayDescr {
///         PyArrayDescr::object(py)
///     }
/// }
///
/// Python::with_gil(|py| {
///     let array = Array2::<Wrapper>::from_shape_fn((2, 3), |(_i, _j)| {
///         Wrapper(Py::new(py, CustomElement).unwrap())
///     });
///
///     let _array: &PyArray<Wrapper, _> = array.to_pyarray(py);
/// });
/// ```
pub unsafe trait Element: Clone + Send {
    /// Flag that indicates whether this type is trivially copyable.
    ///
    /// It should be set to true for all trivially copyable types (like scalar types
    /// and record/array types only containing trivially copyable fields and elements).
    ///
    /// This flag should *always* be set to `false` for object types or record types
    /// that contain object-type fields.
    const IS_COPY: bool;

    /// Returns the associated array descriptor ("dtype") for the given type.
    fn get_dtype(py: Python) -> &PyArrayDescr;
}

fn npy_int_type_lookup<T, T0, T1, T2>(npy_types: [NPY_TYPES; 3]) -> NPY_TYPES {
    // `npy_common.h` defines the integer aliases. In order, it checks:
    // NPY_BITSOF_LONG, NPY_BITSOF_LONGLONG, NPY_BITSOF_INT, NPY_BITSOF_SHORT, NPY_BITSOF_CHAR
    // and assigns the alias to the first matching size, so we should check in this order.
    match size_of::<T>() {
        x if x == size_of::<T0>() => npy_types[0],
        x if x == size_of::<T1>() => npy_types[1],
        x if x == size_of::<T2>() => npy_types[2],
        _ => panic!("Unable to match integer type descriptor: {:?}", npy_types),
    }
}

fn npy_int_type<T: Bounded + Zero + Sized + PartialEq>() -> NPY_TYPES {
    let is_unsigned = T::min_value() == T::zero();
    let bit_width = size_of::<T>() << 3;

    match (is_unsigned, bit_width) {
        (false, 8) => NPY_TYPES::NPY_BYTE,
        (false, 16) => NPY_TYPES::NPY_SHORT,
        (false, 32) => npy_int_type_lookup::<i32, c_long, c_int, c_short>([
            NPY_TYPES::NPY_LONG,
            NPY_TYPES::NPY_INT,
            NPY_TYPES::NPY_SHORT,
        ]),
        (false, 64) => npy_int_type_lookup::<i64, c_long, c_longlong, c_int>([
            NPY_TYPES::NPY_LONG,
            NPY_TYPES::NPY_LONGLONG,
            NPY_TYPES::NPY_INT,
        ]),
        (true, 8) => NPY_TYPES::NPY_UBYTE,
        (true, 16) => NPY_TYPES::NPY_USHORT,
        (true, 32) => npy_int_type_lookup::<u32, c_ulong, c_uint, c_ushort>([
            NPY_TYPES::NPY_ULONG,
            NPY_TYPES::NPY_UINT,
            NPY_TYPES::NPY_USHORT,
        ]),
        (true, 64) => npy_int_type_lookup::<u64, c_ulong, c_ulonglong, c_uint>([
            NPY_TYPES::NPY_ULONG,
            NPY_TYPES::NPY_ULONGLONG,
            NPY_TYPES::NPY_UINT,
        ]),
        _ => unreachable!(),
    }
}

macro_rules! impl_element_scalar {
    (@impl: $ty:ty, $npy_type:expr $(,#[$meta:meta])*) => {
        $(#[$meta])*
        unsafe impl Element for $ty {
            const IS_COPY: bool = true;
            fn get_dtype(py: Python) -> &PyArrayDescr {
                PyArrayDescr::from_npy_type(py, $npy_type)
            }
        }
    };
    ($ty:ty => $npy_type:ident $(,#[$meta:meta])*) => {
        impl_element_scalar!(@impl: $ty, NPY_TYPES::$npy_type $(,#[$meta])*);
    };
    ($($tys:ty),+) => {
        $(impl_element_scalar!(@impl: $tys, npy_int_type::<$tys>());)+
    };
}

impl_element_scalar!(bool => NPY_BOOL);
impl_element_scalar!(i8, i16, i32, i64);
impl_element_scalar!(u8, u16, u32, u64);
impl_element_scalar!(f32 => NPY_FLOAT);
impl_element_scalar!(f64 => NPY_DOUBLE);
impl_element_scalar!(Complex32 => NPY_CFLOAT,
    #[doc = "Complex type with `f32` components which maps to `np.csingle` (`np.complex64`)."]);
impl_element_scalar!(Complex64 => NPY_CDOUBLE,
    #[doc = "Complex type with `f64` components which maps to `np.cdouble` (`np.complex128`)."]);

#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
impl_element_scalar!(usize, isize);

unsafe impl Element for PyObject {
    const IS_COPY: bool = false;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        PyArrayDescr::object(py)
    }
}

#[cfg(test)]
mod tests {
    use std::mem::size_of;

    use super::{dtype, Complex32, Complex64, Element};

    #[test]
    fn test_dtype_names() {
        fn type_name<T: Element>(py: pyo3::Python) -> &str {
            dtype::<T>(py).get_type().name().unwrap()
        }
        pyo3::Python::with_gil(|py| {
            assert_eq!(type_name::<bool>(py), "bool_");
            assert_eq!(type_name::<i8>(py), "int8");
            assert_eq!(type_name::<i16>(py), "int16");
            assert_eq!(type_name::<i32>(py), "int32");
            assert_eq!(type_name::<i64>(py), "int64");
            assert_eq!(type_name::<u8>(py), "uint8");
            assert_eq!(type_name::<u16>(py), "uint16");
            assert_eq!(type_name::<u32>(py), "uint32");
            assert_eq!(type_name::<u64>(py), "uint64");
            assert_eq!(type_name::<f32>(py), "float32");
            assert_eq!(type_name::<f64>(py), "float64");
            assert_eq!(type_name::<Complex32>(py), "complex64");
            assert_eq!(type_name::<Complex64>(py), "complex128");
            match size_of::<usize>() {
                32 => {
                    assert_eq!(type_name::<usize>(py), "uint32");
                    assert_eq!(type_name::<isize>(py), "int32");
                }
                64 => {
                    assert_eq!(type_name::<usize>(py), "uint64");
                    assert_eq!(type_name::<isize>(py), "int64");
                }
                _ => {}
            }
        });
    }
}
<commit_msg>(Use `cfg` instead of `size_of::<usize>` in tests)<commit_after>use std::mem::size_of;
use std::os::raw::{c_int, c_long, c_longlong, c_short, c_uint, c_ulong, c_ulonglong, c_ushort};

use num_traits::{Bounded, Zero};
use pyo3::{ffi, prelude::*, pyobject_native_type_core, types::PyType, AsPyPointer, PyNativeType};

use crate::npyffi::{NpyTypes, PyArray_Descr, NPY_TYPES, PY_ARRAY_API};

pub use num_complex::{Complex32, Complex64};

/// Binding of [`numpy.dtype`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.html).
///
/// # Example
/// ```
/// use pyo3::types::IntoPyDict;
/// pyo3::Python::with_gil(|py| {
///     let locals = [("np", numpy::get_array_module(py).unwrap())].into_py_dict(py);
///     let dtype: &numpy::PyArrayDescr = py
///         .eval("np.array([1, 2, 3.0]).dtype", Some(locals), None)
///         .unwrap()
///         .downcast()
///         .unwrap();
///     assert!(dtype.is_equiv_to(numpy::dtype::<f64>(py)));
/// });
/// ```
pub struct PyArrayDescr(PyAny);

pyobject_native_type_core!(
    PyArrayDescr,
    *PY_ARRAY_API.get_type_object(NpyTypes::PyArrayDescr_Type),
    #module=Some("numpy"),
    #checkfunction=arraydescr_check
);

unsafe fn arraydescr_check(op: *mut ffi::PyObject) -> c_int {
    ffi::PyObject_TypeCheck(
        op,
        PY_ARRAY_API.get_type_object(NpyTypes::PyArrayDescr_Type),
    )
}

/// Returns the type descriptor ("dtype") for a registered type.
pub fn dtype<T: Element>(py: Python) -> &PyArrayDescr {
    T::get_dtype(py)
}

impl PyArrayDescr {
    /// Returns `self` as `*mut PyArray_Descr`.
    pub fn as_dtype_ptr(&self) -> *mut PyArray_Descr {
        self.as_ptr() as _
    }

    /// Returns `self` as `*mut PyArray_Descr` while increasing the reference count.
    ///
    /// Useful in cases where the descriptor is stolen by the API.
    pub fn into_dtype_ptr(&self) -> *mut PyArray_Descr {
        self.into_ptr() as _
    }

    /// Returns the internal `PyType` that this `dtype` holds.
    ///
    /// # Example
    /// ```
    /// pyo3::Python::with_gil(|py| {
    ///    let array = numpy::PyArray::from_vec(py, vec![0.0, 1.0, 2.0f64]);
    ///    let dtype = array.dtype();
    ///    assert_eq!(dtype.get_type().name().unwrap().to_string(), "float64");
    /// });
    /// ```
    pub fn get_type(&self) -> &PyType {
        let dtype_type_ptr = unsafe { *self.as_dtype_ptr() }.typeobj;
        unsafe { PyType::from_type_ptr(self.py(), dtype_type_ptr) }
    }

    /// Shortcut for creating a descriptor of 'object' type.
    pub fn object(py: Python) -> &Self {
        Self::from_npy_type(py, NPY_TYPES::NPY_OBJECT)
    }

    /// Returns the type descriptor ("dtype") for a registered type.
    pub fn of<T: Element>(py: Python) -> &Self {
        T::get_dtype(py)
    }

    /// Returns true if two type descriptors are equivalent.
    pub fn is_equiv_to(&self, other: &Self) -> bool {
        unsafe { PY_ARRAY_API.PyArray_EquivTypes(self.as_dtype_ptr(), other.as_dtype_ptr()) != 0 }
    }

    fn from_npy_type(py: Python, npy_type: NPY_TYPES) -> &Self {
        unsafe {
            let descr = PY_ARRAY_API.PyArray_DescrFromType(npy_type as _);
            py.from_owned_ptr(descr as _)
        }
    }

    /// Retrieves the
    /// [enumerated type](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types)
    /// for this type descriptor.
    pub fn get_typenum(&self) -> c_int {
        unsafe { *self.as_dtype_ptr() }.type_num
    }
}

/// Represents that a type can be an element of `PyArray`.
///
/// Currently, only integer/float/complex types are supported.
/// If you come up with a nice implementation for some other types, we're happy to receive your PR :)
/// You may refer to the [numpy document](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types)
/// for all types that numpy supports.
///
/// # Safety
///
/// A type `T` that implements this trait should be safe when managed in numpy
/// array, thus implementing this trait is marked unsafe. Data types that don't
/// contain Python objects (i.e., either the object type itself or record types
/// containing object-type fields) are assumed to be trivially copyable, which
/// is reflected in the `IS_COPY` flag. Furthermore, it is assumed that for
/// the object type the elements are pointers into the Python heap and that the
/// corresponding `Clone` implemenation will never panic as it only increases
/// the reference count.
///
/// # Custom element types
///
/// You can implement this trait to manage arrays of custom element types, but they still need to be stored
/// on Python's heap using PyO3's [Py](pyo3::Py) type.
///
/// ```
/// use numpy::{ndarray::Array2, Element, PyArray, PyArrayDescr, ToPyArray};
/// use pyo3::{pyclass, Py, Python};
///
/// #[pyclass]
/// pub struct CustomElement;
///
/// // The transparent wrapper is necessary as one cannot implement
/// // a foreign trait (`Element`) on a foreign type (`Py`) directly.
/// #[derive(Clone)]
/// #[repr(transparent)]
/// pub struct Wrapper(pub Py<CustomElement>);
///
/// unsafe impl Element for Wrapper {
///     const IS_COPY: bool = false;
///
///     fn get_dtype(py: Python) -> &PyArrayDescr {
///         PyArrayDescr::object(py)
///     }
/// }
///
/// Python::with_gil(|py| {
///     let array = Array2::<Wrapper>::from_shape_fn((2, 3), |(_i, _j)| {
///         Wrapper(Py::new(py, CustomElement).unwrap())
///     });
///
///     let _array: &PyArray<Wrapper, _> = array.to_pyarray(py);
/// });
/// ```
pub unsafe trait Element: Clone + Send {
    /// Flag that indicates whether this type is trivially copyable.
    ///
    /// It should be set to true for all trivially copyable types (like scalar types
    /// and record/array types only containing trivially copyable fields and elements).
    ///
    /// This flag should *always* be set to `false` for object types or record types
    /// that contain object-type fields.
    const IS_COPY: bool;

    /// Returns the associated array descriptor ("dtype") for the given type.
    fn get_dtype(py: Python) -> &PyArrayDescr;
}

fn npy_int_type_lookup<T, T0, T1, T2>(npy_types: [NPY_TYPES; 3]) -> NPY_TYPES {
    // `npy_common.h` defines the integer aliases. In order, it checks:
    // NPY_BITSOF_LONG, NPY_BITSOF_LONGLONG, NPY_BITSOF_INT, NPY_BITSOF_SHORT, NPY_BITSOF_CHAR
    // and assigns the alias to the first matching size, so we should check in this order.
    match size_of::<T>() {
        x if x == size_of::<T0>() => npy_types[0],
        x if x == size_of::<T1>() => npy_types[1],
        x if x == size_of::<T2>() => npy_types[2],
        _ => panic!("Unable to match integer type descriptor: {:?}", npy_types),
    }
}

fn npy_int_type<T: Bounded + Zero + Sized + PartialEq>() -> NPY_TYPES {
    let is_unsigned = T::min_value() == T::zero();
    let bit_width = size_of::<T>() << 3;

    match (is_unsigned, bit_width) {
        (false, 8) => NPY_TYPES::NPY_BYTE,
        (false, 16) => NPY_TYPES::NPY_SHORT,
        (false, 32) => npy_int_type_lookup::<i32, c_long, c_int, c_short>([
            NPY_TYPES::NPY_LONG,
            NPY_TYPES::NPY_INT,
            NPY_TYPES::NPY_SHORT,
        ]),
        (false, 64) => npy_int_type_lookup::<i64, c_long, c_longlong, c_int>([
            NPY_TYPES::NPY_LONG,
            NPY_TYPES::NPY_LONGLONG,
            NPY_TYPES::NPY_INT,
        ]),
        (true, 8) => NPY_TYPES::NPY_UBYTE,
        (true, 16) => NPY_TYPES::NPY_USHORT,
        (true, 32) => npy_int_type_lookup::<u32, c_ulong, c_uint, c_ushort>([
            NPY_TYPES::NPY_ULONG,
            NPY_TYPES::NPY_UINT,
            NPY_TYPES::NPY_USHORT,
        ]),
        (true, 64) => npy_int_type_lookup::<u64, c_ulong, c_ulonglong, c_uint>([
            NPY_TYPES::NPY_ULONG,
            NPY_TYPES::NPY_ULONGLONG,
            NPY_TYPES::NPY_UINT,
        ]),
        _ => unreachable!(),
    }
}

macro_rules! impl_element_scalar {
    (@impl: $ty:ty, $npy_type:expr $(,#[$meta:meta])*) => {
        $(#[$meta])*
        unsafe impl Element for $ty {
            const IS_COPY: bool = true;
            fn get_dtype(py: Python) -> &PyArrayDescr {
                PyArrayDescr::from_npy_type(py, $npy_type)
            }
        }
    };
    ($ty:ty => $npy_type:ident $(,#[$meta:meta])*) => {
        impl_element_scalar!(@impl: $ty, NPY_TYPES::$npy_type $(,#[$meta])*);
    };
    ($($tys:ty),+) => {
        $(impl_element_scalar!(@impl: $tys, npy_int_type::<$tys>());)+
    };
}

impl_element_scalar!(bool => NPY_BOOL);
impl_element_scalar!(i8, i16, i32, i64);
impl_element_scalar!(u8, u16, u32, u64);
impl_element_scalar!(f32 => NPY_FLOAT);
impl_element_scalar!(f64 => NPY_DOUBLE);
impl_element_scalar!(Complex32 => NPY_CFLOAT,
    #[doc = "Complex type with `f32` components which maps to `np.csingle` (`np.complex64`)."]);
impl_element_scalar!(Complex64 => NPY_CDOUBLE,
    #[doc = "Complex type with `f64` components which maps to `np.cdouble` (`np.complex128`)."]);

#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
impl_element_scalar!(usize, isize);

unsafe impl Element for PyObject {
    const IS_COPY: bool = false;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        PyArrayDescr::object(py)
    }
}

#[cfg(test)]
mod tests {
    use super::{dtype, Complex32, Complex64, Element};

    #[test]
    fn test_dtype_names() {
        fn type_name<T: Element>(py: pyo3::Python) -> &str {
            dtype::<T>(py).get_type().name().unwrap()
        }
        pyo3::Python::with_gil(|py| {
            assert_eq!(type_name::<bool>(py), "bool_");
            assert_eq!(type_name::<i8>(py), "int8");
            assert_eq!(type_name::<i16>(py), "int16");
            assert_eq!(type_name::<i32>(py), "int32");
            assert_eq!(type_name::<i64>(py), "int64");
            assert_eq!(type_name::<u8>(py), "uint8");
            assert_eq!(type_name::<u16>(py), "uint16");
            assert_eq!(type_name::<u32>(py), "uint32");
            assert_eq!(type_name::<u64>(py), "uint64");
            assert_eq!(type_name::<f32>(py), "float32");
            assert_eq!(type_name::<f64>(py), "float64");
            assert_eq!(type_name::<Complex32>(py), "complex64");
            assert_eq!(type_name::<Complex64>(py), "complex128");
            #[cfg(target_pointer_width = "32")]
            {
                assert_eq!(type_name::<usize>(py), "uint32");
                assert_eq!(type_name::<isize>(py), "int32");
            }
            #[cfg(target_pointer_width = "64")]
            {
                assert_eq!(type_name::<usize>(py), "uint64");
                assert_eq!(type_name::<isize>(py), "int64");
            }
        });
    }
}
<|endoftext|>


<commit_before>extern crate tvdb;
extern crate argparse;

use argparse::{ArgumentParser, StoreTrue, Store};

fn main(){
    let mut series_name = "".to_owned();
    let mut season_no = 1;
    let mut episode_no = 1;

    {
        let mut ap = ArgumentParser::new();
        ap.refer(&mut series_name).add_argument("series", Store, "Series name");
        ap.refer(&mut season_no).add_argument("season", Store, "Season number");
        ap.refer(&mut episode_no).add_argument("episode", Store, "Episode number");
        ap.parse_args_or_exit();
    }

    // Construct API object
    let api = tvdb::Tvdb::new("0629B785CE550C8D");

    // Search for series
    let sr = api.search(&series_name, "en").ok().unwrap();

    // Loop over found series
    for r in sr.iter(){
        // Print: "Series Name" (id: 12345)
        println!("{:?} (id: {})", r.seriesname, r.seriesid);

        // Get episode information
        println!("{:?}", api.episode(&sr[0], season_no, episode_no));
    }
}
<commit_msg>Remove unused StoreTrue from example<commit_after>extern crate tvdb;
extern crate argparse;

use argparse::{ArgumentParser, Store};

fn main(){
    let mut series_name = "".to_owned();
    let mut season_no = 1;
    let mut episode_no = 1;

    {
        let mut ap = ArgumentParser::new();
        ap.refer(&mut series_name).add_argument("series", Store, "Series name");
        ap.refer(&mut season_no).add_argument("season", Store, "Season number");
        ap.refer(&mut episode_no).add_argument("episode", Store, "Episode number");
        ap.parse_args_or_exit();
    }

    // Construct API object
    let api = tvdb::Tvdb::new("0629B785CE550C8D");

    // Search for series
    let sr = api.search(&series_name, "en").ok().unwrap();

    // Loop over found series
    for r in sr.iter(){
        // Print: "Series Name" (id: 12345)
        println!("{:?} (id: {})", r.seriesname, r.seriesid);

        // Get episode information
        println!("{:?}", api.episode(&sr[0], season_no, episode_no));
    }
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use std::iter;

use hir::{Adt, HasSource, Semantics};
use ra_ide_db::RootDatabase;
use ra_syntax::ast::{self, edit::IndentLevel, make, AstNode, NameOwner};

use crate::{Assist, AssistCtx, AssistId};

// Assist: fill_match_arms
//
// Adds missing clauses to a `match` expression.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         <|>
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } => (),
//         Action::Stop => (),
//     }
// }
// ```
pub(crate) fn fill_match_arms(ctx: AssistCtx) -> Option<Assist> {
    let match_expr = ctx.find_node_at_offset::<ast::MatchExpr>()?;
    let match_arm_list = match_expr.match_arm_list()?;

    // We already have some match arms, so we don't provide any assists.
    // Unless if there is only one trivial match arm possibly created
    // by match postfix complete. Trivial match arm is the catch all arm.
    let mut existing_arms = match_arm_list.arms();
    if let Some(arm) = existing_arms.next() {
        if !is_trivial(&arm) || existing_arms.next().is_some() {
            return None;
        }
    };

    let expr = match_expr.expr()?;
    let enum_def = resolve_enum_def(&ctx.sema, &expr)?;
    let module = ctx.sema.scope(expr.syntax()).module()?;

    let variants = enum_def.variants(ctx.db);
    if variants.is_empty() {
        return None;
    }

    let db = ctx.db;

    ctx.add_assist(AssistId("fill_match_arms"), "Fill match arms", |edit| {
        let indent_level = IndentLevel::from_node(match_arm_list.syntax());

        let new_arm_list = {
            let arms = variants
                .into_iter()
                .filter_map(|variant| build_pat(db, module, variant))
                .map(|pat| make::match_arm(iter::once(pat), make::expr_unit()));
            indent_level.increase_indent(make::match_arm_list(arms))
        };

        edit.target(match_expr.syntax().text_range());
        edit.set_cursor(expr.syntax().text_range().start());
        edit.replace_ast(match_arm_list, new_arm_list);
    })
}

fn is_trivial(arm: &ast::MatchArm) -> bool {
    match arm.pat() {
        Some(ast::Pat::PlaceholderPat(..)) => true,
        _ => false,
    }
}

fn resolve_enum_def(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<hir::Enum> {
    sema.type_of_expr(&expr)?.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
        Some(Adt::Enum(e)) => Some(e),
        _ => None,
    })
}

fn build_pat(db: &RootDatabase, module: hir::Module, var: hir::EnumVariant) -> Option<ast::Pat> {
    let path = crate::ast_transform::path_to_ast(module.find_use_path(db, var.into())?);

    // FIXME: use HIR for this; it doesn't currently expose struct vs. tuple vs. unit variants though
    let pat: ast::Pat = match var.source(db).value.kind() {
        ast::StructKind::Tuple(field_list) => {
            let pats =
                iter::repeat(make::placeholder_pat().into()).take(field_list.fields().count());
            make::tuple_struct_pat(path, pats).into()
        }
        ast::StructKind::Record(field_list) => {
            let pats = field_list.fields().map(|f| make::bind_pat(f.name().unwrap()).into());
            make::record_pat(path, pats).into()
        }
        ast::StructKind::Unit => make::path_pat(path),
    };

    Some(pat)
}

#[cfg(test)]
mod tests {
    use crate::helpers::{check_assist, check_assist_target};

    use super::fill_match_arms;

    #[test]
    fn fill_match_arms_empty_body() {
        check_assist(
            fill_match_arms,
            r#"
            enum A {
                As,
                Bs,
                Cs(String),
                Ds(String, String),
                Es{ x: usize, y: usize }
            }

            fn main() {
                let a = A::As;
                match a<|> {}
            }
            "#,
            r#"
            enum A {
                As,
                Bs,
                Cs(String),
                Ds(String, String),
                Es{ x: usize, y: usize }
            }

            fn main() {
                let a = A::As;
                match <|>a {
                    A::As => (),
                    A::Bs => (),
                    A::Cs(_) => (),
                    A::Ds(_, _) => (),
                    A::Es { x, y } => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn test_fill_match_arm_refs() {
        check_assist(
            fill_match_arms,
            r#"
            enum A {
                As,
            }

            fn foo(a: &A) {
                match a<|> {
                }
            }
            "#,
            r#"
            enum A {
                As,
            }

            fn foo(a: &A) {
                match <|>a {
                    A::As => (),
                }
            }
            "#,
        );

        check_assist(
            fill_match_arms,
            r#"
            enum A {
                Es{ x: usize, y: usize }
            }

            fn foo(a: &mut A) {
                match a<|> {
                }
            }
            "#,
            r#"
            enum A {
                Es{ x: usize, y: usize }
            }

            fn foo(a: &mut A) {
                match <|>a {
                    A::Es { x, y } => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn fill_match_arms_target() {
        check_assist_target(
            fill_match_arms,
            r#"
            enum E { X, Y }

            fn main() {
                match E::X<|> {}
            }
            "#,
            "match E::X {}",
        );
    }

    #[test]
    fn fill_match_arms_trivial_arm() {
        check_assist(
            fill_match_arms,
            r#"
            enum E { X, Y }

            fn main() {
                match E::X {
                    <|>_ => {},
                }
            }
            "#,
            r#"
            enum E { X, Y }

            fn main() {
                match <|>E::X {
                    E::X => (),
                    E::Y => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn fill_match_arms_qualifies_path() {
        check_assist(
            fill_match_arms,
            r#"
            mod foo { pub enum E { X, Y } }
            use foo::E::X;

            fn main() {
                match X {
                    <|>
                }
            }
            "#,
            r#"
            mod foo { pub enum E { X, Y } }
            use foo::E::X;

            fn main() {
                match <|>X {
                    X => (),
                    foo::E::Y => (),
                }
            }
            "#,
        );
    }
}
<commit_msg>'Fill match arms' should work with existing match arms<commit_after>//! FIXME: write short doc here

use std::{collections::LinkedList, iter};

use hir::{Adt, HasSource, Semantics};
use ra_ide_db::RootDatabase;

use crate::{Assist, AssistCtx, AssistId};
use ra_syntax::{
    ast::{self, edit::IndentLevel, make, AstNode, NameOwner},
    SyntaxKind, SyntaxNode,
};

use ast::{MatchArm, MatchGuard, Pat};

// Assist: fill_match_arms
//
// Adds missing clauses to a `match` expression.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         <|>
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move { distance } => (),
//         Action::Stop => (),
//     }
// }
// ```
pub(crate) fn fill_match_arms(ctx: AssistCtx) -> Option<Assist> {
    let match_expr = ctx.find_node_at_offset::<ast::MatchExpr>()?;
    let match_arm_list = match_expr.match_arm_list()?;

    let expr = match_expr.expr()?;
    let enum_def = resolve_enum_def(&ctx.sema, &expr)?;
    let module = ctx.sema.scope(expr.syntax()).module()?;

    let variants = enum_def.variants(ctx.db);
    if variants.is_empty() {
        return None;
    }

    let db = ctx.db;
    ctx.add_assist(AssistId("fill_match_arms"), "Fill match arms", |edit| {
        let mut arms: Vec<MatchArm> = match_arm_list.arms().collect();
        if arms.len() == 1 {
            if let Some(Pat::PlaceholderPat(..)) = arms[0].pat() {
                arms.clear();
            }
        }

        let mut has_partial_match = false;
        let variants: Vec<MatchArm> = variants
            .into_iter()
            .filter_map(|variant| build_pat(db, module, variant))
            .filter(|variant_pat| {
                !arms.iter().filter_map(|arm| arm.pat().map(|_| arm)).any(|arm| {
                    let pat = arm.pat().unwrap();

                    // Special casee OrPat as separate top-level pats
                    let pats: Vec<Pat> = match Pat::from(pat.clone()) {
                        Pat::OrPat(pats) => pats.pats().collect::<Vec<_>>(),
                        _ => vec![pat],
                    };

                    pats.iter().any(|pat| {
                        match does_arm_pat_match_variant(pat, arm.guard(), variant_pat) {
                            ArmMatch::Yes => true,
                            ArmMatch::No => false,
                            ArmMatch::Partial => {
                                has_partial_match = true;
                                true
                            }
                        }
                    })
                })
            })
            .map(|pat| make::match_arm(iter::once(pat), make::expr_unit()))
            .collect();

        arms.extend(variants);
        if has_partial_match {
            arms.push(make::match_arm(
                iter::once(make::placeholder_pat().into()),
                make::expr_unit(),
            ));
        }

        let indent_level = IndentLevel::from_node(match_arm_list.syntax());
        let new_arm_list = indent_level.increase_indent(make::match_arm_list(arms));

        edit.target(match_expr.syntax().text_range());
        edit.set_cursor(expr.syntax().text_range().start());
        edit.replace_ast(match_arm_list, new_arm_list);
    })
}

enum ArmMatch {
    Yes,
    No,
    Partial,
}

fn does_arm_pat_match_variant(arm: &Pat, arm_guard: Option<MatchGuard>, var: &Pat) -> ArmMatch {
    let arm = flatten_pats(arm.clone());
    let var = flatten_pats(var.clone());
    let mut arm = arm.iter();
    let mut var = var.iter();

    // If the first part of the Pat don't match, there's no match
    match (arm.next(), var.next()) {
        (Some(arm), Some(var)) if arm.text() == var.text() => {}
        _ => return ArmMatch::No,
    }

    // If we have a guard we automatically know we have a partial match
    if arm_guard.is_some() {
        return ArmMatch::Partial;
    }

    if arm.clone().count() != var.clone().count() {
        return ArmMatch::Partial;
    }

    let direct_match = arm.zip(var).all(|(arm, var)| {
        if arm.text() == var.text() {
            return true;
        }
        match (arm.kind(), var.kind()) {
            (SyntaxKind::PLACEHOLDER_PAT, SyntaxKind::PLACEHOLDER_PAT) => true,
            (SyntaxKind::DOT_DOT_PAT, SyntaxKind::PLACEHOLDER_PAT) => true,
            (SyntaxKind::BIND_PAT, SyntaxKind::PLACEHOLDER_PAT) => true,
            _ => false,
        }
    });

    match direct_match {
        true => ArmMatch::Yes,
        false => ArmMatch::Partial,
    }
}

fn flatten_pats(pat: Pat) -> Vec<SyntaxNode> {
    let mut pats: LinkedList<SyntaxNode> = pat.syntax().children().collect();
    let mut out: Vec<SyntaxNode> = vec![];
    while let Some(p) = pats.pop_front() {
        pats.extend(p.children());
        out.push(p);
    }
    out
}

fn resolve_enum_def(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<hir::Enum> {
    sema.type_of_expr(&expr)?.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
        Some(Adt::Enum(e)) => Some(e),
        _ => None,
    })
}

fn build_pat(db: &RootDatabase, module: hir::Module, var: hir::EnumVariant) -> Option<ast::Pat> {
    let path = crate::ast_transform::path_to_ast(module.find_use_path(db, var.into())?);

    // FIXME: use HIR for this; it doesn't currently expose struct vs. tuple vs. unit variants though
    let pat: ast::Pat = match var.source(db).value.kind() {
        ast::StructKind::Tuple(field_list) => {
            let pats =
                iter::repeat(make::placeholder_pat().into()).take(field_list.fields().count());
            make::tuple_struct_pat(path, pats).into()
        }
        ast::StructKind::Record(field_list) => {
            let pats = field_list.fields().map(|f| make::bind_pat(f.name().unwrap()).into());
            make::record_pat(path, pats).into()
        }
        ast::StructKind::Unit => make::path_pat(path),
    };

    Some(pat)
}

#[cfg(test)]
mod tests {
    use crate::helpers::{check_assist, check_assist_target};

    use super::fill_match_arms;

    #[test]
    fn partial_fill_multi() {
        check_assist(
            fill_match_arms,
            r#"
            enum A {
                As,
                Bs(i32, Option<i32>)
            }
            fn main() {
                match A::As<|> {
                    A::Bs(_, Some(_)) => (),
                }
            }
            "#,
            r#"
            enum A {
                As,
                Bs(i32, Option<i32>)
            }
            fn main() {
                match <|>A::As {
                    A::Bs(_, Some(_)) => (),
                    A::As => (),
                    _ => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn partial_fill_record() {
        check_assist(
            fill_match_arms,
            r#"
            enum A {
                As,
                Bs{x:i32, y:Option<i32>},
            }
            fn main() {
                match A::As<|> {
                    A::Bs{x,y:Some(_)} => (),
                }
            }
            "#,
            r#"
            enum A {
                As,
                Bs{x:i32, y:Option<i32>},
            }
            fn main() {
                match <|>A::As {
                    A::Bs{x,y:Some(_)} => (),
                    A::As => (),
                    _ => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn partial_fill_or_pat() {
        check_assist(
            fill_match_arms,
            r#"
            enum A {
                As,
                Bs,
                Cs(Option<i32>),
            }
            fn main() {
                match A::As<|> {
                    A::Cs(_) | A::Bs => (),
                }
            }
            "#,
            r#"
            enum A {
                As,
                Bs,
                Cs(Option<i32>),
            }
            fn main() {
                match <|>A::As {
                    A::Cs(_) | A::Bs => (),
                    A::As => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn partial_fill_or_pat2() {
        check_assist(
            fill_match_arms,
            r#"
            enum A {
                As,
                Bs,
                Cs(Option<i32>),
            }
            fn main() {
                match A::As<|> {
                    A::Cs(Some(_)) | A::Bs => (),
                }
            }
            "#,
            r#"
            enum A {
                As,
                Bs,
                Cs(Option<i32>),
            }
            fn main() {
                match <|>A::As {
                    A::Cs(Some(_)) | A::Bs => (),
                    A::As => (),
                    _ => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn partial_fill() {
        check_assist(
            fill_match_arms,
            r#"
            enum A {
                As,
                Bs,
                Cs,
                Ds(String),
                Es(B),
            }
            enum B {
                Xs,
                Ys,
            }
            fn main() {
                match A::As<|> {
                    A::Bs if 0 < 1 => (),
                    A::Ds(_value) => (),
                    A::Es(B::Xs) => (),
                }
            }
            "#,
            r#"
            enum A {
                As,
                Bs,
                Cs,
                Ds(String),
                Es(B),
            }
            enum B {
                Xs,
                Ys,
            }
            fn main() {
                match <|>A::As {
                    A::Bs if 0 < 1 => (),
                    A::Ds(_value) => (),
                    A::Es(B::Xs) => (),
                    A::As => (),
                    A::Cs => (),
                    _ => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn fill_match_arms_empty_body() {
        check_assist(
            fill_match_arms,
            r#"
            enum A {
                As,
                Bs,
                Cs(String),
                Ds(String, String),
                Es{ x: usize, y: usize }
            }

            fn main() {
                let a = A::As;
                match a<|> {}
            }
            "#,
            r#"
            enum A {
                As,
                Bs,
                Cs(String),
                Ds(String, String),
                Es{ x: usize, y: usize }
            }

            fn main() {
                let a = A::As;
                match <|>a {
                    A::As => (),
                    A::Bs => (),
                    A::Cs(_) => (),
                    A::Ds(_, _) => (),
                    A::Es { x, y } => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn test_fill_match_arm_refs() {
        check_assist(
            fill_match_arms,
            r#"
            enum A {
                As,
            }

            fn foo(a: &A) {
                match a<|> {
                }
            }
            "#,
            r#"
            enum A {
                As,
            }

            fn foo(a: &A) {
                match <|>a {
                    A::As => (),
                }
            }
            "#,
        );

        check_assist(
            fill_match_arms,
            r#"
            enum A {
                Es{ x: usize, y: usize }
            }

            fn foo(a: &mut A) {
                match a<|> {
                }
            }
            "#,
            r#"
            enum A {
                Es{ x: usize, y: usize }
            }

            fn foo(a: &mut A) {
                match <|>a {
                    A::Es { x, y } => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn fill_match_arms_target() {
        check_assist_target(
            fill_match_arms,
            r#"
            enum E { X, Y }

            fn main() {
                match E::X<|> {}
            }
            "#,
            "match E::X {}",
        );
    }

    #[test]
    fn fill_match_arms_trivial_arm() {
        check_assist(
            fill_match_arms,
            r#"
            enum E { X, Y }

            fn main() {
                match E::X {
                    <|>_ => {},
                }
            }
            "#,
            r#"
            enum E { X, Y }

            fn main() {
                match <|>E::X {
                    E::X => (),
                    E::Y => (),
                }
            }
            "#,
        );
    }

    #[test]
    fn fill_match_arms_qualifies_path() {
        check_assist(
            fill_match_arms,
            r#"
            mod foo { pub enum E { X, Y } }
            use foo::E::X;

            fn main() {
                match X {
                    <|>
                }
            }
            "#,
            r#"
            mod foo { pub enum E { X, Y } }
            use foo::E::X;

            fn main() {
                match <|>X {
                    X => (),
                    foo::E::Y => (),
                }
            }
            "#,
        );
    }
}
<|endoftext|>


<commit_before>extern crate syn;

use syn::DeriveInput;

pub fn generate_elm(_: &DeriveInput) {
    // do nothing
}
<commit_msg>Write a test file to generated directory<commit_after>extern crate syn;

use std::env;
use std::fs::File;
use std::io::{Write, BufWriter};
use syn::DeriveInput;

pub fn generate_elm(_: &DeriveInput) {
    let mut path = env::current_dir().unwrap();
    path.push("generated/test.txt");

    let file = File::create(path).unwrap();
    let ref mut w = BufWriter::new(file);
    w.write("Hello, World!".as_bytes()).unwrap();
}
<|endoftext|>