//
// Cforall Version 1.0.0 Copyright (C) 2021 University of Waterloo
//
// The contents of this file are covered under the licence agreement in the
// file "LICENCE" distributed with Cforall.
//
// channel.hfa -- LIBCFATHREAD
// Runtime locks that used with the runtime thread system.
//
// Author           : Colby Alexander Parsons
// Created On       : Thu Jan 21 19:46:50 2022
// Last Modified By :
// Last Modified On :
// Update Count     :
//

#pragma once

#include <locks.hfa>
#include <list.hfa>
#include "select.hfa"

// returns true if woken due to shutdown
// blocks thread on list and releases passed lock
static inline bool block( dlist( select_node ) & queue, void * elem_ptr, go_mutex & lock ) {
    select_node sn{ active_thread(), elem_ptr };
    insert_last( queue, sn );
    unlock( lock );
    park();
    __atomic_thread_fence( __ATOMIC_SEQ_CST );
    return sn.extra == 0p;
}

// Waituntil support (un)register_select helper routine
// Sets select node avail if not special OR case and then unlocks
static inline void __set_avail_then_unlock( select_node & node, go_mutex & mutex_lock ) {
    if ( node.park_counter ) __make_select_node_available( node );
    unlock( mutex_lock );
}

// void * used for some fields since exceptions don't work with parametric polymorphism currently
exception channel_closed {
    // on failed insert elem is a ptr to the element attempting to be inserted
    // on failed remove elem ptr is 0p
    // on resumption of a failed insert this elem will be inserted
    // so a user may modify it in the resumption handler
    void * elem;

    // pointer to chan that is closed
    void * closed_chan;
};
vtable(channel_closed) channel_closed_vt;

static inline bool is_insert( channel_closed & e ) { return e.elem != 0p; }
static inline bool is_remove( channel_closed & e ) { return e.elem == 0p; }

// #define CHAN_STATS // define this to get channel stats printed in dtor

forall( T ) {

struct __attribute__((aligned(128))) channel {
    size_t size, front, back, count;
    T * buffer;
    dlist( select_node ) prods, cons; // lists of blocked threads
    go_mutex mutex_lock;              // MX lock
    bool closed;                      // indicates channel close/open
    #ifdef CHAN_STATS
    size_t p_blocks, p_ops, c_blocks, c_ops;      // counts total ops and ops resulting in a blocked thd
    #endif
};
static inline void ?{}( channel(T) & this, channel(T) this2 ) = void;
static inline void ?=?( channel(T) & this, channel(T) this2 ) = void;

static inline void ?{}( channel(T) &c, size_t _size ) with(c) {
    size = _size;
    front = back = count = 0;
    if ( size != 0 ) buffer = aalloc( size );
    prods{};
    cons{};
    mutex_lock{};
    closed = false;
    #ifdef CHAN_STATS
    p_blocks = 0;
    p_ops = 0;
    c_blocks = 0;
    c_ops = 0;
    #endif
}

static inline void ?{}( channel(T) &c ){ ((channel(T) &)c){ 0 }; }
static inline void ^?{}( channel(T) &c ) with(c) {
    #ifdef CHAN_STATS
    printf("Channel %p Blocks: %lu,\t\tOperations: %lu,\t%.2f%% of ops blocked\n", &c, p_blocks + c_blocks, p_ops + c_ops, ((double)p_blocks + c_blocks)/(p_ops + c_ops) * 100);
    printf("Channel %p Consumer Blocks: %lu,\tConsumer Ops: %lu,\t%.2f%% of Consumer ops blocked\n", &c, p_blocks, p_ops, ((double)p_blocks)/p_ops * 100);
    printf("Channel %p Producer Blocks: %lu,\tProducer Ops: %lu,\t%.2f%% of Producer ops blocked\n", &c, c_blocks, c_ops, ((double)c_blocks)/c_ops * 100);
    #endif
    verifyf( __handle_waituntil_OR( cons ) || __handle_waituntil_OR( prods ) || cons`isEmpty && prods`isEmpty, 
        "Attempted to delete channel with waiting threads (Deadlock).\n" );
    if ( size != 0 ) delete( buffer );
}
static inline size_t get_count( channel(T) & chan ) with(chan) { return __atomic_load_n( &count, __ATOMIC_RELAXED ); }
static inline size_t get_size( channel(T) & chan ) with(chan) { return __atomic_load_n( &size, __ATOMIC_RELAXED ); }
static inline bool has_waiters( channel(T) & chan ) with(chan) { return !cons`isEmpty || !prods`isEmpty; }
static inline bool has_waiting_consumers( channel(T) & chan ) with(chan) { return !cons`isEmpty; }
static inline bool has_waiting_producers( channel(T) & chan ) with(chan) { return !prods`isEmpty; }

// closes the channel and notifies all blocked threads
static inline void close( channel(T) & chan ) with(chan) {
    lock( mutex_lock );
    closed = true;

    // flush waiting consumers and producers
    while ( has_waiting_consumers( chan ) ) {
        if( !__handle_waituntil_OR( cons ) ) // ensure we only signal special OR case threads when they win the race
            break;  // if __handle_waituntil_OR returns false cons is empty so break
        cons`first.extra = 0p;
        wake_one( cons );
    }
    while ( has_waiting_producers( chan ) ) {
        if( !__handle_waituntil_OR( prods ) ) // ensure we only signal special OR case threads when they win the race
            break;  // if __handle_waituntil_OR returns false prods is empty so break
        prods`first.extra = 0p;
        wake_one( prods );
    }
    unlock(mutex_lock);
}

static inline void is_closed( channel(T) & chan ) with(chan) { return closed; }

// used to hand an element to a blocked consumer and signal it
static inline void __cons_handoff( channel(T) & chan, T & elem ) with(chan) {
    memcpy( cons`first.extra, (void *)&elem, sizeof(T) ); // do waiting consumer work
    __atomic_thread_fence( __ATOMIC_SEQ_CST );
    wake_one( cons );
}

// used to hand an element to a blocked producer and signal it
static inline void __prods_handoff( channel(T) & chan, T & retval ) with(chan) {
    memcpy( (void *)&retval, prods`first.extra, sizeof(T) );
    __atomic_thread_fence( __ATOMIC_SEQ_CST );
    wake_one( prods );
}

static inline void flush( channel(T) & chan, T elem ) with(chan) {
    lock( mutex_lock );
    while ( count == 0 && !cons`isEmpty ) {
        __cons_handoff( chan, elem );
    }
    unlock( mutex_lock );
}

// handles buffer insert
static inline void __buf_insert( channel(T) & chan, T & elem ) with(chan) {
    memcpy( (void *)&buffer[back], (void *)&elem, sizeof(T) );
    count += 1;
    back++;
    if ( back == size ) back = 0;
}

// needed to avoid an extra copy in closed case
static inline bool __internal_try_insert( channel(T) & chan, T & elem ) with(chan) {
    lock( mutex_lock );
    #ifdef CHAN_STATS
    p_ops++;
    #endif

    ConsEmpty: if ( !cons`isEmpty ) {
        if ( !__handle_waituntil_OR( cons ) ) break ConsEmpty;
        __cons_handoff( chan, elem );
        unlock( mutex_lock );
        return true;
    }

    if ( count == size ) { unlock( mutex_lock ); return false; }

    __buf_insert( chan, elem );
    unlock( mutex_lock );
    return true;
}

// attempts a nonblocking insert
// returns true if insert was successful, false otherwise
static inline bool try_insert( channel(T) & chan, T elem ) { return __internal_try_insert( chan, elem ); }

// handles closed case of insert routine
static inline void __closed_insert( channel(T) & chan, T & elem ) with(chan) {
    channel_closed except{ &channel_closed_vt, &elem, &chan };
    throwResume except; // throw closed resumption
    if ( !__internal_try_insert( chan, elem ) ) throw except; // if try to insert fails (would block), throw termination
}

static inline void insert( channel(T) & chan, T elem ) with(chan) {
    // check for close before acquire mx
    if ( unlikely(closed) ) {
        __closed_insert( chan, elem );
        return;
    } 

    lock( mutex_lock );

    #ifdef CHAN_STATS
    if ( !closed ) p_ops++;
    #endif

    // if closed handle
    if ( unlikely(closed) ) {
        unlock( mutex_lock );
        __closed_insert( chan, elem );
        return;
    }

    // buffer count must be zero if cons are blocked (also handles zero-size case)
    ConsEmpty: if ( !cons`isEmpty ) {
        if ( !__handle_waituntil_OR( cons ) ) break ConsEmpty;
        __cons_handoff( chan, elem );
        unlock( mutex_lock );
        return;
    }

    // wait if buffer is full, work will be completed by someone else
    if ( count == size ) {
        #ifdef CHAN_STATS
        p_blocks++;
        #endif

        // check for if woken due to close
        if ( unlikely( block( prods, &elem, mutex_lock ) ) )
            __closed_insert( chan, elem );
        return;
    } // if

    __buf_insert( chan, elem );
    unlock( mutex_lock );
}

// does the buffer remove and potentially does waiting producer work
static inline void __do_remove( channel(T) & chan, T & retval ) with(chan) {
    memcpy( (void *)&retval, (void *)&buffer[front], sizeof(T) );
    count -= 1;
    front = (front + 1) % size;
    if (count == size - 1 && !prods`isEmpty ) {
        if ( !__handle_waituntil_OR( prods ) ) return;
        __buf_insert( chan, *(T *)prods`first.extra );  // do waiting producer work
        wake_one( prods );
    }
}

// needed to avoid an extra copy in closed case and single return val case
static inline bool __internal_try_remove( channel(T) & chan, T & retval ) with(chan) {
    lock( mutex_lock );
    #ifdef CHAN_STATS
    c_ops++;
    #endif

    ZeroSize: if ( size == 0 && !prods`isEmpty ) {
        if ( !__handle_waituntil_OR( prods ) ) break ZeroSize;
        __prods_handoff( chan, retval );
        unlock( mutex_lock );
        return true;
    }

    if ( count == 0 ) { unlock( mutex_lock ); return false; }

    __do_remove( chan, retval );
    unlock( mutex_lock );
    return true;
}

// attempts a nonblocking remove
// returns [T, true] if insert was successful
// returns [T, false] if insert was successful (T uninit)
static inline [T, bool] try_remove( channel(T) & chan ) {
    T retval;
    bool success = __internal_try_remove( chan, retval );
    return [ retval, success ];
}

static inline T try_remove( channel(T) & chan ) {
    T retval;
    __internal_try_remove( chan, retval );
    return retval;
}

// handles closed case of insert routine
static inline void __closed_remove( channel(T) & chan, T & retval ) with(chan) {
    channel_closed except{ &channel_closed_vt, 0p, &chan };
    throwResume except; // throw resumption
    if ( !__internal_try_remove( chan, retval ) ) throw except; // if try to remove fails (would block), throw termination
}

static inline T remove( channel(T) & chan ) with(chan) {
    T retval;
    if ( unlikely(closed) ) {
        __closed_remove( chan, retval );
        return retval;
    } 
    lock( mutex_lock );

    #ifdef CHAN_STATS
    if ( !closed ) c_ops++;
    #endif

    if ( unlikely(closed) ) {
        unlock( mutex_lock );
        __closed_remove( chan, retval );
        return retval;
    } 

    // have to check for the zero size channel case
    ZeroSize: if ( size == 0 && !prods`isEmpty ) {
        if ( !__handle_waituntil_OR( prods ) ) break ZeroSize;
        __prods_handoff( chan, retval );
        unlock( mutex_lock );
        return retval;
    }

    // wait if buffer is empty, work will be completed by someone else
    if ( count == 0 ) {
        #ifdef CHAN_STATS
        c_blocks++;
        #endif
        // check for if woken due to close
        if ( unlikely( block( cons, &retval, mutex_lock ) ) )
            __closed_remove( chan, retval );
        return retval;
    }

    // Remove from buffer
    __do_remove( chan, retval );
    unlock( mutex_lock );
    return retval;
}
static inline void remove( channel(T) & chan ) { T elem = (T)remove( chan ); }


///////////////////////////////////////////////////////////////////////////////////////////
// The following is Go-style operator support for channels
///////////////////////////////////////////////////////////////////////////////////////////

static inline void ?<<?( channel(T) & chan, T elem ) { insert( chan, elem ); }
static inline void ?<<?( T & ret, channel(T) & chan ) { ret = remove( chan ); }

///////////////////////////////////////////////////////////////////////////////////////////
// The following is support for waituntil (select) statements
///////////////////////////////////////////////////////////////////////////////////////////
static inline bool unregister_chan( channel(T) & chan, select_node & node ) with(chan) {
    if ( !node`isListed && !node.park_counter ) return false; // handle special OR case
    lock( mutex_lock );
    if ( node`isListed ) { // op wasn't performed
        remove( node );
        unlock( mutex_lock );
        return false;
    }
    unlock( mutex_lock );

    // only return true when not special OR case and status is SAT
    return !node.park_counter ? false : *node.clause_status == __SELECT_SAT;
}

// special case of __handle_waituntil_OR, that does some work to avoid starvation/deadlock case
static inline bool __handle_pending( dlist( select_node ) & queue, select_node & mine ) {
    while ( !queue`isEmpty ) {
        // if node not a special OR case or if we win the special OR case race break
        if ( !queue`first.clause_status || queue`first.park_counter || __pending_set_other( queue`first, mine, ((unsigned long int)(&(queue`first))) ) )
            return true;
        
        // our node lost the race when toggling in __pending_set_other
        if ( *mine.clause_status != __SELECT_PENDING ) 
            return false;

        // otherwise we lost the special OR race so discard node
        try_pop_front( queue );
    }
    return false;
}

// type used by select statement to capture a chan read as the selected operation
struct chan_read {
    T * ret;
    channel(T) * chan;
};
__CFA_SELECT_GET_TYPE( chan_read(T) );

static inline void ?{}( chan_read(T) & cr, channel(T) * chan, T * ret ) {
    cr.chan = chan;
    cr.ret = ret;
}
static inline chan_read(T) ?<<?( T & ret, channel(T) & chan ) { chan_read(T) cr{ &chan, &ret }; return cr; }

static inline void __handle_select_closed_read( chan_read(T) & this, select_node & node ) with(*this.chan, this) {
    __closed_remove( *chan, *ret );
    // if we get here then the insert succeeded
    __make_select_node_available( node );
}

static inline bool register_select( chan_read(T) & this, select_node & node ) with(*this.chan, this) {
    lock( mutex_lock );
    node.extra = ret; // set .extra so that if it == 0p later in on_selected it is due to channel close

    #ifdef CHAN_STATS
    if ( !closed ) c_ops++;
    #endif

    if ( !node.park_counter ) {
        // are we special case OR and front of cons is also special case OR
        if ( !unlikely(closed) && !prods`isEmpty && prods`first.clause_status && !prods`first.park_counter ) {
            if ( !__make_select_node_pending( node ) ) {
                unlock( mutex_lock );
                return false;
            }

            if ( __handle_pending( prods, node ) ) {
                __prods_handoff( *chan, *ret );
                __make_select_node_sat( node ); // need to to mark SAT now that we know operation is done or else threads could get stuck in __mark_select_node
                unlock( mutex_lock );
                return true;
            }
            if ( *node.clause_status == __SELECT_PENDING )
                __make_select_node_unsat( node );
        }
        // check if we can complete operation. If so race to establish winner in special OR case
        if ( count != 0 || !prods`isEmpty || unlikely(closed) ) {
            if ( !__make_select_node_available( node ) ) { // we didn't win the race so give up on registering
                unlock( mutex_lock );
                return false;
            }
        }
    }

    if ( unlikely(closed) ) {
        unlock( mutex_lock );
        __handle_select_closed_read( this, node );
        return true;
    }

    // have to check for the zero size channel case
    ZeroSize: if ( size == 0 && !prods`isEmpty ) {
        if ( !__handle_waituntil_OR( prods ) ) break ZeroSize;
        __prods_handoff( *chan, *ret );
        __set_avail_then_unlock( node, mutex_lock );
        return true;
    }

    // wait if buffer is empty, work will be completed by someone else
    if ( count == 0 ) {
        #ifdef CHAN_STATS
        c_blocks++;
        #endif
        
        insert_last( cons, node );
        unlock( mutex_lock );
        return false;
    }

    // Remove from buffer
    __do_remove( *chan, *ret );
    __set_avail_then_unlock( node, mutex_lock );
    return true;
}
static inline bool unregister_select( chan_read(T) & this, select_node & node ) { return unregister_chan( *this.chan, node ); }
static inline bool on_selected( chan_read(T) & this, select_node & node ) with(this) {
    if ( unlikely(node.extra == 0p) ) {
        if ( !exception_in_flight() ) __closed_remove( *chan, *ret ); // check if woken up due to closed channel
        else return false;
    }
    // This is only reachable if not closed or closed exception was handled
    return true;
}

// type used by select statement to capture a chan read as the selected operation that doesn't have a param to read to
struct chan_read_no_ret {
    T retval;
    chan_read( T ) c_read;
};
__CFA_SELECT_GET_TYPE( chan_read_no_ret(T) );

static inline void ?{}( chan_read_no_ret(T) & this, channel(T) & chan ) {
    this.c_read{ &chan, &this.retval };
}

static inline chan_read_no_ret(T) remove( channel(T) & chan ) { chan_read_no_ret(T) c_read{ chan }; return c_read; }
static inline bool register_select( chan_read_no_ret(T) & this, select_node & node ) { 
    this.c_read.ret = &this.retval;
    return register_select( this.c_read, node );
}
static inline bool unregister_select( chan_read_no_ret(T) & this, select_node & node ) { return unregister_select( this.c_read, node ); }
static inline bool on_selected( chan_read_no_ret(T) & this, select_node & node ) { return on_selected( this.c_read, node ); }

// type used by select statement to capture a chan write as the selected operation
struct chan_write {
    T elem;
    channel(T) * chan;
};
__CFA_SELECT_GET_TYPE( chan_write(T) );

static inline void ?{}( chan_write(T) & cw, channel(T) * chan, T elem ) {
    cw.chan = chan;
    memcpy( (void *)&cw.elem, (void *)&elem, sizeof(T) );
}
static inline chan_write(T) ?<<?( channel(T) & chan, T elem ) { chan_write(T) cw{ &chan, elem }; return cw; }
static inline chan_write(T) insert( T elem, channel(T) & chan) { chan_write(T) cw{ &chan, elem }; return cw; }

static inline void __handle_select_closed_write( chan_write(T) & this, select_node & node ) with(*this.chan, this) {
    __closed_insert( *chan, elem );
    // if we get here then the insert succeeded
    __make_select_node_available( node );
}

static inline bool register_select( chan_write(T) & this, select_node & node ) with(*this.chan, this) {
    lock( mutex_lock );
    node.extra = &elem; // set .extra so that if it == 0p later in on_selected it is due to channel close

    #ifdef CHAN_STATS
    if ( !closed ) p_ops++;
    #endif

    // special OR case handling
    if ( !node.park_counter ) {
        // are we special case OR and front of cons is also special case OR
        if ( !unlikely(closed) && !cons`isEmpty && cons`first.clause_status && !cons`first.park_counter ) {
            if ( !__make_select_node_pending( node ) ) {
                unlock( mutex_lock );
                return false;
            }

            if ( __handle_pending( cons, node ) ) {
                __cons_handoff( *chan, elem );
                __make_select_node_sat( node ); // need to to mark SAT now that we know operation is done or else threads could get stuck in __mark_select_node
                unlock( mutex_lock );
                return true;
            }
            if ( *node.clause_status == __SELECT_PENDING )
                __make_select_node_unsat( node );
        }
        // check if we can complete operation. If so race to establish winner in special OR case
        if ( count != size || !cons`isEmpty || unlikely(closed) ) {
            if ( !__make_select_node_available( node ) ) { // we didn't win the race so give up on registering
                unlock( mutex_lock );
                return false;
            }
        }
    }

    // if closed handle
    if ( unlikely(closed) ) {
        unlock( mutex_lock );
        __handle_select_closed_write( this, node );
        return true;
    }

    // handle blocked consumer case via handoff (buffer is implicitly empty)
    ConsEmpty: if ( !cons`isEmpty ) {
        if ( !__handle_waituntil_OR( cons ) ) break ConsEmpty;
        __cons_handoff( *chan, elem );
        __set_avail_then_unlock( node, mutex_lock );
        return true;
    }

    // insert node in list if buffer is full, work will be completed by someone else
    if ( count == size ) {
        #ifdef CHAN_STATS
        p_blocks++;
        #endif

        insert_last( prods, node );
        unlock( mutex_lock );
        return false;
    } // if

    // otherwise carry out write either via normal insert
    __buf_insert( *chan, elem );
    __set_avail_then_unlock( node, mutex_lock );
    return true;
}
static inline bool unregister_select( chan_write(T) & this, select_node & node ) { return unregister_chan( *this.chan, node ); }

static inline bool on_selected( chan_write(T) & this, select_node & node ) with(this) { 
    if ( unlikely(node.extra == 0p) ) {
        if ( !exception_in_flight() ) __closed_insert( *chan, elem ); // check if woken up due to closed channel
        else return false;
    }
    // This is only reachable if not closed or closed exception was handled
    return true;
}

} // forall( T )