1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

use ndarray::prelude::*;
use ndarray::{Ix2, Data, DataMut};
use lapack::c::Layout;
use std::slice;

/// Return an array with the specified dimensions and layout.
///
/// This function is used internally to ensure that
pub fn matrix_with_layout<T: Default>(d: Ix2, layout: Layout) -> Array<T, Ix2> {
    Array::default(match layout {
        Layout::RowMajor => d.into(),
        Layout::ColumnMajor => d.f(),
    })
}

/// Return the slice, layout, and leading dimension of the matrix.
///
/// For LAPACKE methods, the memory layout does not need to be
/// contiguous. We only required that either rows or columns are
/// contiguous.
pub fn slice_and_layout_mut<D, S: DataMut<Elem = D>>(mat: &mut ArrayBase<S, Ix2>)
                                                     -> Option<(&mut [S::Elem], Layout, Ixs)> {
    let strides = {
        let s = mat.strides();
        (s[0], s[1])
    };

    if strides.0 < 0 || strides.1 < 0 {
        return None;
    }

    let dim = mat.dim();

    // One of the stides, must be 1
    if strides.1 == 1 {
        let m = strides.0;
        let s = unsafe {
            let nelem: usize = (dim.0 - 1) * m as usize + dim.1;
            slice::from_raw_parts_mut(mat.as_mut_ptr(), nelem)
        };
        Some((s, Layout::RowMajor, m))
    } else if strides.0 == 1 {
        let n = strides.1;
        let s = unsafe {
            let nelem: usize = (dim.1 - 1) * n as usize + dim.0;
            slice::from_raw_parts_mut(mat.as_mut_ptr(), nelem)
        };
        Some((s, Layout::ColumnMajor, n))
    } else {
        None
    }
}

/// Return the slice, layout, and leading dimension of the matrix.
///
/// For LAPACKE methods, the memory layout does not need to be
/// contiguous. We only required that either rows or columns are
/// contiguous.
pub fn slice_and_layout<D, S: Data<Elem = D>>(mat: &ArrayBase<S, Ix2>)
                                              -> Option<(&[S::Elem], Layout, Ixs)> {
    let strides = {
        let s = mat.strides();
        (s[0], s[1])
    };

    if strides.0 < 0 || strides.1 < 0 {
        return None;
    }

    let dim = mat.dim();

    // One of the stides, must be 1
    if strides.1 == 1 {
        let m = strides.0;
        let s = unsafe {
            let nelem: usize = (dim.0 - 1) * m as usize + dim.1;
            slice::from_raw_parts(mat.as_ptr(), nelem)
        };
        Some((s, Layout::RowMajor, m))
    } else if strides.0 == 1 {
        let n = strides.1;
        let s = unsafe {
            let nelem: usize = (dim.1 - 1) * n as usize + dim.0;
            slice::from_raw_parts(mat.as_ptr(), nelem)
        };
        Some((s, Layout::ColumnMajor, n))
    } else {
        None
    }
}


/// Return the slice, layout, and leading dimension of the matrix. If
/// the matrix can be interpreted with multiple inputs, assume the
/// previous one.
///
/// For LAPACKE methods, the memory layout does not need to be
/// contiguous. We only required that either rows or columns are
/// contiguous.
///
/// Returns None if the layout can't be matched.
pub fn slice_and_layout_matching_mut<D, S: DataMut<Elem = D>>(mat: &mut ArrayBase<S, Ix2>,
                                                              layout: Layout)
                                                              -> Option<(&mut [S::Elem], Ixs)> {
    let dim = mat.dim();

    // For column vectors, we can choose whatever layout we want.
    if dim.1 == 1 {
        let m = mat.strides()[0];

        let s = unsafe {
            let nelem: usize = (dim.0 - 1) * m as usize + dim.1;
            slice::from_raw_parts_mut(mat.as_mut_ptr(), nelem)
        };
        return Some((s, m));
    }

    // Otherwise, we just use the normal method and check for a match.
    if let Some((slice, lo, ld)) = slice_and_layout_mut(mat) {
        if lo == layout {
            return Some((slice, ld));
        }
    }

    None
}