use std::thread::LocalKey;
use std::cell::RefCell;
use crate::IndexT;
use crate::tape::Tape;
use crate::tape::sealed::ThisThreadTape;
use crate::op::id;
pub fn doc_generic_v() {}
#[derive(Clone, Debug)]
pub struct AD<V> {
pub(crate) tape_id : usize,
pub(crate) var_index : usize,
pub(crate) value : V,
}
impl<V> AD<V> {
pub(crate) fn new(
new_tape_id: usize, new_var_index: usize, new_value: V )-> Self {
Self {
tape_id : new_tape_id,
var_index : new_var_index,
value : new_value,
}
}
}
impl<V> AD<V> {
pub fn to_value(self) -> V {
self.value
}
}
impl<V : std::fmt::Display> std::fmt::Display for AD<V> {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "{}", self.value)
}
}
pub fn doc_ad_binary_op() { }
macro_rules! ad_binary_op { ($Name:ident, $Op:tt) => { paste::paste! {
fn [< record_ $Name:lower _aa >]<V> (
tape: &mut Tape<V> ,
lhs: &AD<V> ,
rhs: &AD<V> ,
) -> (usize, usize)
where
V : Clone ,
{
let mut new_tape_id = 0;
let mut new_var_index = 0;
if tape.recording {
let var_lhs = lhs.tape_id == tape.tape_id;
let var_rhs = rhs.tape_id == tape.tape_id;
if var_lhs || var_rhs {
new_tape_id = tape.tape_id;
new_var_index = tape.n_var;
tape.n_var += 1;
tape.op2arg.push( tape.arg_all.len() as IndexT );
if var_lhs && var_rhs {
tape.id_all.push( id::[< $Name:upper _VV_OP >] );
tape.arg_all.push( lhs.var_index as IndexT );
tape.arg_all.push( rhs.var_index as IndexT );
} else if var_lhs {
tape.id_all.push( id::[< $Name:upper _VC_OP >] );
tape.arg_all.push( lhs.var_index as IndexT );
tape.arg_all.push( tape.con_all.len() as IndexT );
tape.con_all.push( rhs.value.clone() );
} else {
tape.id_all.push( id::[< $Name:upper _CV_OP >] );
tape.arg_all.push( tape.con_all.len() as IndexT );
tape.con_all.push( lhs.value.clone() );
tape.arg_all.push( rhs.var_index as IndexT );
}
}
}
( new_tape_id, new_var_index )
}
#[doc = concat!(
"& `AD<V>` ", stringify!($Op), " & `AD<V>`",
"; see [doc_ad_binary_op]"
)]
impl<V> std::ops::$Name< &AD<V> > for &AD<V>
where
for<'a> &'a V: std::ops::$Name<&'a V, Output=V>,
V : Clone + crate::ThisThreadTapePublic ,
{ type Output = AD<V>;
fn [< $Name:lower >](self , rhs : &AD<V> ) -> AD<V>
{
let new_value = &self.value $Op &rhs.value;
let local_key : &LocalKey< RefCell< Tape<V> > > =
ThisThreadTape::get();
let (new_tape_id, new_var_index) =
local_key.with_borrow_mut( |tape|
[< record_ $Name:lower _aa >]::<V> ( tape, &self, &rhs )
);
AD::new(new_tape_id, new_var_index, new_value)
}
}
fn [< record_ $Name:lower _av >]<V> (
tape: &mut Tape<V> ,
lhs: &AD<V> ,
rhs: &V ,
) -> (usize, usize)
where
V : Clone ,
{
let mut new_tape_id = 0;
let mut new_var_index = 0;
if tape.recording {
let var_lhs = lhs.tape_id == tape.tape_id;
if var_lhs {
new_tape_id = tape.tape_id;
new_var_index = tape.n_var;
tape.n_var += 1;
tape.op2arg.push( tape.arg_all.len() as IndexT );
tape.id_all.push( id::[< $Name:upper _VC_OP >] );
tape.arg_all.push( lhs.var_index as IndexT );
tape.arg_all.push( tape.con_all.len() as IndexT );
tape.con_all.push( rhs.clone() );
}
}
(new_tape_id, new_var_index)
}
#[doc = concat!(
"& `AD<V>` ", stringify!($Op), " & V`",
"; see [doc_ad_binary_op]"
)]
impl<V> std::ops::$Name< &V> for &AD<V>
where
for<'a> &'a V: std::ops::$Name<&'a V, Output=V>,
V : Clone + crate::ThisThreadTapePublic ,
{ type Output = AD<V>;
fn [< $Name:lower >](self , rhs : &V ) -> AD<V>
{
let new_value = &self.value $Op &rhs;
let local_key : &LocalKey< RefCell< Tape<V> > > =
ThisThreadTape::get();
let (new_tape_id, new_var_index) =
local_key.with_borrow_mut( |tape|
[< record_ $Name:lower _av >]::<V> ( tape, &self, &rhs )
);
AD::new(new_tape_id, new_var_index, new_value)
}
}
} } }
ad_binary_op!(Add, +);
ad_binary_op!(Sub, -);
ad_binary_op!(Mul, *);
ad_binary_op!(Div, /);
pub fn doc_ad_compound_op() { }
macro_rules! ad_compound_op { ($Name:ident, $Op:tt) => { paste::paste! {
fn [< record_ $Name:lower _assign_aa >]<V> (
tape: &mut Tape<V> ,
lhs: &mut AD<V> ,
rhs: & AD<V> )
where
V : Clone,
{
if tape.recording {
let var_lhs = lhs.tape_id == tape.tape_id;
let var_rhs = rhs.tape_id == tape.tape_id;
if var_lhs || var_rhs {
tape.op2arg.push( tape.arg_all.len() as IndexT );
if var_lhs && var_rhs {
tape.id_all.push( id::[< $Name:upper _VV_OP >] );
tape.arg_all.push( lhs.var_index as IndexT);
tape.arg_all.push( rhs.var_index as IndexT);
} else if var_lhs {
tape.id_all.push( id::[< $Name:upper _VC_OP >] );
tape.arg_all.push( lhs.var_index as IndexT);
tape.arg_all.push( tape.con_all.len() as IndexT );
tape.con_all.push( rhs.value.clone() );
} else {
tape.id_all.push( id::[< $Name:upper _CV_OP >] );
tape.arg_all.push( tape.con_all.len() as IndexT );
tape.con_all.push( lhs.value.clone() );
tape.arg_all.push( rhs.var_index as IndexT);
}
lhs.tape_id = tape.tape_id;
lhs.var_index = tape.n_var;
tape.n_var += 1;
}
}
}
#[doc = concat!(
"`AD<V>` ", stringify!($Op), " & `AD<V>`",
"; see [doc_ad_compound_op]"
)]
impl<V> std::ops::[< $Name Assign >] < &AD<V> > for AD<V>
where
V: Clone +
for<'a> std::ops::[< $Name Assign >] <&'a V> +
crate::ThisThreadTapePublic ,
{ fn [< $Name:lower _assign >] (&mut self, rhs : &AD<V> )
{ let local_key : &LocalKey< RefCell< Tape<V> > > =
ThisThreadTape::get();
local_key.with_borrow_mut( |tape|
[< record_ $Name:lower _assign_aa >]::<V> ( tape, self, rhs )
);
self.value $Op &rhs.value;
}
}
fn [< record_ $Name:lower _assign_av >]<V> (
tape: &mut Tape<V> ,
lhs: &mut AD<V> ,
rhs: & V )
where
V : Clone,
{
if tape.recording {
let var_lhs = lhs.tape_id == tape.tape_id;
if var_lhs {
tape.op2arg.push( tape.arg_all.len() as IndexT );
tape.id_all.push( id::[< $Name:upper _VC_OP >] );
tape.arg_all.push( lhs.var_index as IndexT);
tape.arg_all.push( tape.con_all.len() as IndexT );
tape.con_all.push( rhs.clone() );
lhs.var_index = tape.n_var;
tape.n_var += 1;
}
}
}
#[doc = concat!(
"`AD<V>` ", stringify!($Op), " & V; see [doc_ad_compound_op]"
)]
impl<V> std::ops::[< $Name Assign >] <&V> for AD<V>
where
V: Clone +
for<'a> std::ops::[< $Name Assign >] <&'a V> +
crate::ThisThreadTapePublic ,
{ fn [< $Name:lower _assign >] (&mut self, rhs : &V)
{ let local_key : &LocalKey< RefCell< Tape<V> > > =
ThisThreadTape::get();
local_key.with_borrow_mut( |tape|
[< record_ $Name:lower _assign_av >]::<V> ( tape, self, rhs )
);
self.value $Op &rhs;
}
}
} } }
ad_compound_op!(Add, +=);
ad_compound_op!(Sub, -=);
ad_compound_op!(Mul, *=);
ad_compound_op!(Div, /=);
macro_rules! record_value_op_ad{ ($Name:ident, $Op:tt) => { paste::paste! {
#[doc = concat!( "record one ", stringify!($Name),
" where lhs is a value and rhs is a variable"
) ]
pub(crate) fn [< record_value_ $Name:lower _ad >]<V> (
tape: &mut Tape<V> ,
lhs: &V ,
rhs: &AD<V> ,
) -> (usize, usize)
where
V : Clone ,
{
let mut new_tape_id = 0;
let mut new_var_index = 0;
if tape.recording {
let var_rhs = rhs.tape_id == tape.tape_id;
if var_rhs {
new_tape_id = tape.tape_id;
new_var_index = tape.n_var;
tape.n_var += 1;
tape.op2arg.push( tape.arg_all.len() as IndexT );
tape.id_all.push( id::[< $Name:upper _CV_OP >] );
tape.arg_all.push( tape.con_all.len() as IndexT );
tape.con_all.push( lhs.clone() );
tape.arg_all.push( rhs.var_index as IndexT );
}
}
(new_tape_id, new_var_index)
}
} } }
record_value_op_ad!(Add, +=);
record_value_op_ad!(Sub, -=);
record_value_op_ad!(Mul, *=);
record_value_op_ad!(Div, /=);
macro_rules! impl_value_op_ad{
($V:ty) => {
crate::ad::impl_value_op_ad!($V, Add, +);
crate::ad::impl_value_op_ad!($V, Sub, -);
crate::ad::impl_value_op_ad!($V, Mul, *);
crate::ad::impl_value_op_ad!($V, Div, /);
};
($V:ty, $Name:ident, $Op:tt) => { paste::paste! {
#[doc =
"see [doc_ad_binary_op](crate::ad::doc_ad_binary_op)"
]
impl std::ops::$Name< &AD<$V> > for & $V
where
for <'a> &'a $V : std::ops::$Name<&'a $V, Output=$V>,
{ type Output = AD<$V>;
#[ doc = concat!(
"compute & `", stringify!($V), "` ",
stringify!($Op), " & `AD<", stringify!($f1), ">` "
) ]
fn [< $Name:lower >]
(self , rhs : &AD<$V>
) -> AD<$V> {
let new_value = self $Op &rhs.value;
let local_key : &LocalKey<
RefCell< crate::tape::Tape<$V> >
> = crate::tape::sealed::ThisThreadTape::get();
let (new_tape_id, new_var_index) = local_key.with_borrow_mut(
|tape|
crate::ad::[< record_value_ $Name:lower _ad >]::<$V>
( tape, &self, &rhs )
);
AD::new(new_tape_id, new_var_index, new_value)
}
}
} }
}
pub(crate) use impl_value_op_ad;
pub fn ad_from_value<V>(value : V) -> AD<V> {
let tape_id = 0;
let var_index = 0;
AD::new(tape_id, var_index, value)
}
pub fn ad_from_vector<V> ( vec : Vec<V> ) -> Vec< AD<V> > {
assert_ne!( vec.len() , 0 );
let tape_id = 0;
let var_index = 0;
let avec = vec.into_iter().map(
|value| AD::new(tape_id, var_index, value)
).collect();
avec
}
pub fn ad_to_vector<V> ( avec : Vec< AD<V> > ) -> Vec<V> {
assert_ne!( avec.len() , 0 );
let vec = avec.into_iter().map( |ad| ad.value).collect();
vec
}
pub fn doc_impl_ad_from_f32() { }
macro_rules! impl_ad_from_f32{ ($V:ty) => {
impl From<f32> for crate::AD<$V> {
fn from( f32_value : f32 ) -> crate::AD<$V> {
let tape_id = 0;
let var_index = 0;
let value : $V = f32_value.into();
crate::AD::new(tape_id, var_index, value)
}
}
} }
pub(crate) use impl_ad_from_f32;
pub fn doc_impl_ad_from_f64() { }
macro_rules! impl_ad_from_f64{ ($V:ty) => {
impl From<f64> for crate::AD<$V> {
fn from( f64_value : f64 ) -> crate::AD<$V> {
let tape_id = 0;
let var_index = 0;
let value : $V = f64_value.into();
crate::AD::new(tape_id, var_index, value)
}
}
} }
pub(crate) use impl_ad_from_f64;