brick-smtlib

// -*- mode: C++; indent-tabs-mode: nil; c-basic-offset: 4 -*-

/*
 * Utilities for printing SMT-LIBv2 formulas
 */

/*
 * (c) 2017 Vladimír Štill <xstill@fi.muni.cz>
 * (c) 2018 Petr Ročkai <code@fixp.eu>
 * (c) 2019 Henrich Lauko <xlauko@mail.muni.cz>
 */

/* Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE. */

#pragma once

#include <brick-assert>
#include <brick-string>
#include <brick-data>
#include <brick-smt>

#include <map>
#include <set>
#include <memory>
#include <variant>
#include <iomanip>
#include <optional>
#include <functional>
#include <unordered_map>
#include <unordered_set>
#include <string_view>
#include <cmath>

namespace brq
{
    using smtlib_format = std::function< void ( brq::string_builder & ) >;
    struct smtlib_context;

    struct smtlib_node
    {
        smtlib_format _fmt;
        int bw = 0;
        enum type_t { t_bool, t_bv, t_float, t_array } type;

        smtlib_node( int bw, type_t type, smtlib_format &&fmt )
            : _fmt( std::move( fmt ) ), bw( bw ), type( type )
        {}

        smtlib_node( int bw, type_t type, std::string str )
            : _fmt( [=]( brq::string_builder &o ) { o << str; } ), bw( bw ), type( type )
        {}

        constexpr inline bool is_bool() const noexcept { return type == t_bool; }
        constexpr inline bool is_bv() const noexcept { return type == t_bv; }
        constexpr inline bool is_float() const noexcept { return type == t_float; }
        constexpr inline bool is_array() const noexcept { return type == t_array; }
    };

    using smtlib_type = smtlib_node::type_t;

    inline brq::string_builder &operator<<( brq::string_builder &o, const smtlib_node &n )
    {
        n._fmt( o );
        return o;
    }

    inline std::string to_string( const smtlib_node &p )
    {
        brq::string_builder ss;
        ss << p;
        if ( ss.truncated() )
            throw std::bad_alloc();
        else
            return ss.buffer();
    }

    inline std::string to_string( smtlib_type type )
    {
        switch ( type )
        {
            case smtlib_node::t_bool: return "bool";
            case smtlib_node::t_bv: return "bv";
            case smtlib_node::t_float: return "float";
            case smtlib_node::t_array: return "array";
        }
    }

    /* float operations */
    enum class smtlib_rounding
    {
        RNE,    // round nearest, ties to even
        RNA,    // round nearest, ties to away (from zero)
        RTP,    // round toward positive
        RTN,    // round toward negative
        RTZ     // round toward zero
    };

    inline std::string to_string( smtlib_rounding mode )
    {
        switch ( mode )
        {
            case smtlib_rounding::RNE: return "RNE";
            case smtlib_rounding::RNA: return "RNA";
            case smtlib_rounding::RTP: return "RTP";
            case smtlib_rounding::RTN: return "RTN";
            case smtlib_rounding::RTZ: return "RTZ";
            default: UNREACHABLE( "unknown rounding mode", mode );
        };
    }

    inline brq::string_builder &operator<<( brq::string_builder &o, smtlib_rounding mode )
    {
        o << to_string( mode );
        return o;
    }

    static constexpr const char *to_fp( size_t bw )
    {
        switch ( bw )
        {
            case 16: return "(_ to_fp 5 11)";
            case 32: return "(_ to_fp 8 24)";
            case 64: return "(_ to_fp 11 53)";
            default: UNREACHABLE( "Unsupported bitwidth." );
        }
    }

    struct smtlib_context
    {
        using node = smtlib_node;
        using type_t = node::type_t;

        std::unordered_set< std::string > def_set;  // set of definitions
        std::vector< std::pair< std::string, node > > defs;  // names of definitions
        std::unordered_map< std::string, node > vars;

        void clear() { def_set.clear(); defs.clear(); vars.clear(); }

        // alias a node with a name
        node define( std::string name, node def )
        {
            ASSERT( !name.empty() );
            ASSERT( !def_set.count( name ) );
            def_set.insert( name );
            defs.emplace_back( name, def );
            return node( def.bw, def.type, name );
        }

        node variable( node n, std::string name )
        {
            if ( vars.count( name ) )
                ASSERT_EQ( to_string( n ), to_string( vars.find( name )->second ) );
            else
                vars.emplace( name, n );
            return node( n.bw, n.type, name );
        }

        node symbol( int bw, type_t type, std::string name )
        {
            return node( bw, type, name );
        }

        node bitvec( int, brick::data::SmallVector< uint64_t > data );
        node boolT() { return node( 1, node::t_bool, "Bool" ); }

        node typeT( type_t type, int bitwidth )
        {
            if ( type == node::t_bv )
                return bitvecT( bitwidth );
            if ( type == node::t_float )
                return floatT( bitwidth );
            UNREACHABLE( "unsupported type" );
        }

        node bitvecT( int bitwidth )
        {
            ASSERT_LEQ( 0, bitwidth );
            return node( bitwidth, node::t_bv, "(_ BitVec " + std::to_string( bitwidth ) + ")" );
        }

        node floatT( int bitwidth )
        {
            ASSERT_LEQ( 0, bitwidth );
            return node( bitwidth, node::t_float, "(Float" + std::to_string( bitwidth ) + ")" );
        }

        node arrayT( type_t type, int value_bw, int index_bw )
        {
            ASSERT_LEQ( 0, value_bw );
            ASSERT_LEQ( 0, index_bw );

            auto values = typeT( type, value_bw );
            auto indices = bitvecT( index_bw );

            auto fmt = [=]( brq::string_builder &o )
            {
                o << "(Array " << values << " " << indices << ")";
            };

            return node( value_bw, type, fmt );
        }

        node bitvec( int bw, uint64_t val )
        {
            return bitvec( bw, brick::data::SmallVector< uint64_t >{ val } );
        }

        template< typename Int, typename =
                std::enable_if_t< std::is_integral< Int >::value &&
                                  ( sizeof( Int ) <= sizeof( uint64_t ) ) > >
        node bitvec( Int i )
        {
            return bitvec( sizeof( Int ) * 8, brick::data::SmallVector< uint64_t >{ uint64_t( i ) } );
        }

        node floatv( int bw, double val, smtlib_rounding mode = smtlib_rounding::RNE );

        void print( brq::string_builder &o, node n, bool exq = true )  // exq = Existential Quantifier
        {
            auto bind = [&]( std::string type, auto &vec, auto next )
            {
                for ( auto s : vec )
                    o << "\n    (" << type << " ((" << s.first << " " << s.second << ")) ";
                next();
                for ( unsigned i = 0; i < vec.size(); ++i )
                    o << ")";
            };

            auto let = [&]{ bind( "let", defs, [&]{ o << n; } ); }; // local variable binder
            if ( exq )
                bind( "exists", vars, let );
            else
                let();
        }

        void query( brq::string_builder &o, node n )
        {
            for ( auto v : vars )
                o << "(declare-fun " << v.first << " () " << v.second << ")\n";
            o << "(assert ";
            print( o, n, false );
            o << ")\n(check-sat)";
        }

        std::string print( node n, bool exq = true )
        {
            brq::string_builder s;
            print( s, n, exq );
            if ( s.truncated() )
                throw std::bad_alloc();
            else
                return s.buffer();
        }

        std::string query( node n )
        {
            brq::string_builder s;
            query( s, n );
            if ( s.truncated() )
                throw std::bad_alloc();
            else
                return s.buffer();
        }

        using Vector = brick::data::SmallVector< node, 2 >;

        bool is_rounding_mode_required( smt_op op )
        {
            return op == smt_op::fp_add || op == smt_op::fp_sub ||
                   op == smt_op::fp_mul || op == smt_op::fp_div;
        }

        node expr( int bw, smt_op op, const Vector &args, std::optional< smtlib_rounding > mode )
        {
            auto fmt = [=]( brq::string_builder &o )
            {
                o << "(" << smt_name( op );

                if ( is_rounding_mode_required( op ) )
                {
                    ASSERT( mode.has_value() );
                    o << " " << to_string( mode.value() );
                }

                for ( auto &a : args )
                    o << " " << a;
                o << ")";
            };

            auto type = smt_traits( op ).type == smt_op_compare ? node::t_bool : args[ 0 ].type;
            return node( bw, type, fmt );
        }

        node expr( int bw, smt_op op, const Vector &args )
        {
            return expr( bw, op, args, std::nullopt );
        }


        template< smt_op op >
        node unop( int bw, node arg )
        {
            static_assert( smt_arity( op ) == 1 );
            return expr( bw, op, { arg } );
        }

        template< smt_op op >
        node binop( int bw, node a, node b )
        {
            static_assert( smt_arity( op ) == 2 );
            ASSERT_EQ( a.type, b.type );
            return expr( bw, op, { a, b } );
        }

        template< smt_op op >
        node fpbinop( int bw, node a, node b )
        {
            static_assert( smt_arity( op ) == 2 );
            ASSERT( a.type == b.type );
            return expr( bw, op, { a, b }, smtlib_rounding::RNE );
        }

        // extract a subrange from a bitvector
        node extract( int highest, int lowest, node arg )
        {
            ASSERT_LEQ( lowest, highest );
            auto f = [=]( brq::string_builder &o )
            {
                o << "((_ extract " << std::dec << highest << " " << lowest << ") " << arg << ")";
            };
            return node( 1 + highest - lowest, node::t_bv, f );
        };

        // if - then - else
        node ite( node cond, node t, node f )
        {
            ASSERT_EQ( t.bw, f.bw );
            ASSERT( t.type == f.type );
            auto fmt = [=]( brq::string_builder &o )
            {
                o << "(ite " << cond << " " << t << " " << f << ")";
            };
            return node( t.bw, t.type, fmt );
        }

        node select( node array, node offset, int bw )
        {
            auto fmt = [=]( brq::string_builder &o )
            {
                o << "(select " << array << " " << offset << ")";
            };
            return node( bw, array.type, fmt ); // FIXME loaded type
        }

        node store( node array, node offset, node value, int bw )
        {
            auto fmt = [=]( brq::string_builder &o )
            {
                o << "(store " << array << " " << offset << " " << value << ")";
            };
            return node( bw, array.type, fmt );
        }

        static constexpr auto rounding_mode = smtlib_rounding::RNE;

        inline std::string to_fp_bitwidth( size_t bw )
        {
            switch ( bw )
            {
                case 16: return "5 11";
                case 32: return "8 24";
                case 64: return "11 53";
                default: UNREACHABLE( "Unsupported bitwidth." );
            }
        }

        template < bool _unsigned >
        inline std::string fp_to_bv_pref()
        {
            std::string pref = "(_ fp.to_";
            if constexpr ( _unsigned )
                return pref + "ubv ";
            return pref + "sbv ";
        }

        template< bool _unsigned >
        inline node fp_to_bv_cast( int bw, const node & arg )
        {
            auto fmt = [=]( brq::string_builder &o )
            {
                o << fp_to_bv_pref< _unsigned >() << bw << ") " << to_string( rounding_mode )
                  << " " << arg << "(_ BitVec " << bw << ") )";
            };
            return node( bw, node::t_bv, fmt );
        }

        template< bool _unsigned >
        inline node bv_to_fp_cast( int bw, const node & arg )
        {
            auto fmt = [=]( brq::string_builder &o )
            {
                o << "(" << to_fp( bw ) << " " << to_string( rounding_mode )
                  << " " << arg << "(_ FloatingPoint " << to_fp_bitwidth( bw ) << "))";
            };
            return node( bw, node::t_float, fmt );
        }

        node cast( smt_op op, int bw, const node & arg )
        {
            if ( op == smt_op::fp_ext || op == smt_op::fp_trunc )
            {
                const auto mode = smtlib_rounding::RNE;
                auto fmt = [=]( brq::string_builder &o )
                {
                    o << "(" << to_fp( bw ) << " " << to_string( mode ) << " " << arg << " )";
                };
                return node( bw, node::t_float, fmt );
            }

            if ( op == smt_op::fp_toubv )
            {
                return fp_to_bv_cast< true /* unsigned */ >( bw, arg );
            }

            if ( op == smt_op::fp_tosbv )
            {
                return fp_to_bv_cast< false /* signed */ >( bw, arg );
            }

            if ( op == smt_op::bv_utofp )
            {
                return bv_to_fp_cast< true /* unsigned */ >( bw, arg );
            }

            if ( op == smt_op::bv_stofp )
            {
                return bv_to_fp_cast< false /* signed */ >( bw, arg );
            }

            NOT_IMPLEMENTED();
        }
    };

    inline static uint64_t u64mask( int bitwidth )
    {
        if ( bitwidth >= 64 )
            return ~uint64_t( 0 );
        return (uint64_t( 1 ) << bitwidth) - 1;
    }

    inline brq::string_builder to_bitvector( int bitwidth,
                                                    brick::data::SmallVector< uint64_t > data )
    {
        int bw = bitwidth;
        int i = 0;

        while ( bw > 0 && i < int( data.size() ) )
        {
            data[ i ] &= u64mask( bw );
            bw -= 64;
            ++i;
        }

        while ( i < int( data.size() ) )
            data.pop_back();
        while ( int( data.size() ) < (bitwidth + 63) / 64 )
            data.push_back( 0 );

        brq::string_builder o;

        if ( bitwidth % 4 == 0 )
        {
            o << "#x" << std::hex;
            // data are little endian
            for ( int i = data.size() - 1; i >= 0; --i )
            {
                int w = std::min( bitwidth - i * 64, 64 ) / 4;
                o << std::hex << brq::pad( w, '0' ) << data[ i ] << brq::mark;
            }
        }
        else
        {
            o << "#b";
            for ( int i = data.size() - 1; i >= 0; --i ) {
                for ( int j = std::min( bitwidth - 1 - i * 64, 63 ); j >= 0; --j )
                    if ( data[i] & (uint64_t( 1 ) << j) )
                        o << "1";
                    else
                        o << "0";
            }
        }

        return o;
    }

    inline smtlib_node smtlib_context::floatv( int bitwidth, double val, smtlib_rounding mode )
    {
        static_assert( sizeof( double ) == sizeof( uint64_t ) );

        auto sign = [] ( auto v ) { return std::signbit( v ) ? "-" : "+"; };

        switch( std::fpclassify( val ) )
        {
            case FP_INFINITE:
                return node( bitwidth, node::t_float
                    , brq::format( "(_", sign( val ), "oo ", to_fp_bitwidth( bitwidth ), ")" ).buffer() );
            case FP_NAN:
                return node( bitwidth, node::t_float
                    , brq::format( "(_ NaN ", to_fp_bitwidth( bitwidth ), ")" ).buffer() );
            case FP_ZERO:
                return node( bitwidth, node::t_float
                    , brq::format( "(_", sign( val ), "zero ", to_fp_bitwidth( bitwidth ), ")" ).buffer() );
            default:
                return node( bitwidth, node::t_float
                    , brq::format( "(", to_fp( bitwidth ), " ", to_string( mode ), " ", val, " )" ).buffer() );
        }
    }

    inline smtlib_node smtlib_context::bitvec( int bitwidth, brick::data::SmallVector< uint64_t > data )
    {
        auto o = to_bitvector( bitwidth, data );
        if ( o.truncated() )
            throw std::bad_alloc();
        return node( bitwidth, node::t_bv, o.buffer() );
    }
}

namespace t_brq::smtlib
{
    using node = brq::smtlib_node;
    using Op = brq::smt_op;

    struct BVecTest
    {
        TEST(string)
        {
            brq::smtlib_context ctx;

            auto a = ctx.bitvec( 42 );
            ASSERT_EQ( "#x0000002a", to_string( a ) );

            auto b = ctx.bitvec( uint64_t( 42 ) );
            ASSERT_EQ( "#x000000000000002a", to_string( b ) );

            auto c = ctx.bitvec( 4, 0xff );
            ASSERT_EQ( "#xf", to_string( c ) );

            brick::data::SmallVector< uint64_t > data = { ~uint64_t( 0 ), ~uint64_t( 0 ) };
            auto d = ctx.bitvec( 96, data );
            ASSERT_EQ( "#x" + std::string( 24, 'f' ), to_string( d ) );

            auto e = ctx.bitvec( 256, data );
            ASSERT_EQ( "#x" + std::string( 32, '0' ) + std::string( 32, 'f' ), to_string( e ) );

            auto f = ctx.bitvec( 96, 42 );
            ASSERT_EQ( "#x" + std::string( 22, '0' ) + "2a", to_string( f ) );

            auto g = ctx.bitvec( 16, data );
            ASSERT_EQ( "#xffff", to_string( g ) );

            auto h = ctx.bitvec( 7, data );
            ASSERT_EQ( "#b1111111", to_string( h ) );

            auto i = ctx.bitvec( 7, 0x29 );
            ASSERT_EQ( "#b0101001", to_string( i ) );

            auto j = ctx.bitvec( 67, { (uint64_t( 1 ) << 63) | 1, 0x5 } );
            ASSERT_EQ( "#b1011" + std::string( 62, '0' ) + "1", to_string( j ) );

            auto k = ctx.bitvec( 67, { (uint64_t( 1 ) << 33) | 1, 0x5 } );
            ASSERT_EQ( "#b101" + std::string( 30, '0' ) + "1" + std::string( 32, '0' ) + "1",
                       to_string( k ) );
        }
    };

    struct ExtractTest
    {
        TEST(print)
        {
            brq::smtlib_context ctx;

            ASSERT_EQ( "((_ extract 32 0) #x000000000000002a)",
                       to_string( ctx.extract( 32, 0, ctx.bitvec( 42ull ) ) ) );
            ASSERT_EQ( "((_ extract 15 15) (concat #xff #xff))",
                       to_string( ctx.extract( 15, 15,
                                               ctx.binop< Op::bv_concat >( 16,
                                                                           ctx.bitvec( 8, 0xff ),
                                                                           ctx.bitvec( 8, 0xff ) ) ) ) );
        }
    };

    struct CombinationTest
    {
        TEST(simple)
        {
            brq::smtlib_context ctx;

            ASSERT_EQ( to_string( ctx.binop< Op::bv_add >( 32, ctx.bitvec( 42 ), ctx.bitvec( 32 ) ) ),
                       "(bvadd #x0000002a #x00000020)" );
            ASSERT_EQ( to_string( ctx.binop< Op::bv_and >( 32, ctx.symbol( 32, node::t_bv, "a" ),
                                                           ctx.bitvec( 3, 5 ) ) ),
                       "(bvand a #b101)" );
            ASSERT_EQ( to_string(
                           ctx.binop< Op::bv_sle >(
                               16, ctx.extract( 32, 16, ctx.symbol( 32, node::t_bv, "a" ) ),
                               ctx.unop< Op::bv_neg >( 16, ctx.bitvec( short( 42 ) ) ) ) ),
                       "(bvsle ((_ extract 32 16) a) (bvneg #x002a))" );
        }
    };
}

// vim: syntax=cpp tabstop=4 shiftwidth=4 expandtab ft=cpp