<commit_before>/************************************************************
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you 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.
*
*************************************************************/

#include "gtest/gtest.h"
#include "communication/msg.h"
using namespace singa;
TEST(MsgTest, AddrTest) {
  int src_grp = 1, src_worker = 2;
  int dst_grp = 0, dst_server = 1;
  int src_addr = Addr(src_grp, src_worker, 0);
  int dst_addr = Addr(dst_grp, dst_server, 1);
  Msg msg(src_addr, dst_addr);
  msg.set_trgt(123, -1);
  ASSERT_EQ(AddrGrp(msg.src()), src_grp);
  ASSERT_EQ(AddrID(msg.src()), src_worker);
  ASSERT_EQ(AddrType(msg.src()), 0);

  msg.SwapAddr();
  ASSERT_EQ(AddrGrp(msg.src()), dst_grp);
  ASSERT_EQ(AddrID(msg.src()), dst_server);
  ASSERT_EQ(AddrType(msg.src()), 1);
  ASSERT_EQ(msg.trgt_val(), 123);
  ASSERT_EQ(msg.trgt_version(), -1);
}

TEST(MsgTest, AddFrameTest) {
  int buf[5]={1,2,3,4,5};
  Msg msg;
  msg.AddFrame("abcdefg", 7);
  msg.AddFrame(buf, sizeof(int) * 5);

  msg.FirstFrame();
  char* str = msg.FrameStr();
  ASSERT_STREQ(str, "abcdefg");
  delete str;
  ASSERT_EQ(msg.NextFrame(), true);
  int *val = static_cast<int*>(msg.FrameData());
  ASSERT_EQ(val[3], 4);
  ASSERT_EQ(msg.NextFrame(), false);

  msg.FirstFrame();
  str = msg.FrameStr();
  ASSERT_STREQ(str, "abcdefg");
  msg.LastFrame();
  val = static_cast<int*>(msg.FrameData());
  ASSERT_EQ(val[2], 3);
}

TEST(MsgTest, AddFormatFrame) {
  int x = 5;
  Msg msg;
  msg.AddFormatFrame("i", 12);
  msg.AddFormatFrame("f", 10.f);
  msg.AddFormatFrame("s", "abc");
  msg.AddFormatFrame("p", &x);
  msg.AddFormatFrame("isfp", 12, "abc", 10.f, &x);

  msg.FirstFrame();
  int y;
  msg.ParseFormatFrame("i", &y);
  ASSERT_EQ(y, 12);
  ASSERT_EQ(msg.NextFrame(), true);

  float z;
  msg.ParseFormatFrame("f", &z);
  ASSERT_EQ(z, 10.f);
  ASSERT_EQ(msg.NextFrame(), true);

  char buf[10];
  msg.ParseFormatFrame("s", buf);
  ASSERT_STREQ(buf, "abc");
  ASSERT_EQ(msg.NextFrame(), true);

  int *p;
  msg.ParseFormatFrame("p", &p);
  ASSERT_EQ(p, &x);
  ASSERT_EQ(msg.NextFrame(), true);

  msg.ParseFormatFrame("isfp", &y, buf, &z, &p);
  ASSERT_EQ(y, 12);
  ASSERT_STREQ(buf, "abc");
  ASSERT_EQ(z, 10.f);
  ASSERT_EQ(p, &x);
}
<commit_msg>SINGA-79 Fix bug in singatool that can not parse -conf flag<commit_after>/************************************************************
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you 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.
*
*************************************************************/

#include "gtest/gtest.h"
#include "comm/msg.h"
using namespace singa;
TEST(MsgTest, AddrTest) {
  int src_grp = 1, src_worker = 2;
  int dst_grp = 0, dst_server = 1;
  int src_addr = Addr(src_grp, src_worker, 0);
  int dst_addr = Addr(dst_grp, dst_server, 1);
  Msg msg(src_addr, dst_addr);
  msg.set_trgt(123, -1);
  ASSERT_EQ(AddrGrp(msg.src()), src_grp);
  ASSERT_EQ(AddrID(msg.src()), src_worker);
  ASSERT_EQ(AddrType(msg.src()), 0);

  msg.SwapAddr();
  ASSERT_EQ(AddrGrp(msg.src()), dst_grp);
  ASSERT_EQ(AddrID(msg.src()), dst_server);
  ASSERT_EQ(AddrType(msg.src()), 1);
  ASSERT_EQ(msg.trgt_val(), 123);
  ASSERT_EQ(msg.trgt_version(), -1);
}

TEST(MsgTest, AddFrameTest) {
  int buf[5]={1,2,3,4,5};
  Msg msg;
  msg.AddFrame("abcdefg", 7);
  msg.AddFrame(buf, sizeof(int) * 5);

  msg.FirstFrame();
  char* str = msg.FrameStr();
  ASSERT_STREQ(str, "abcdefg");
  delete str;
  ASSERT_EQ(msg.NextFrame(), true);
  int *val = static_cast<int*>(msg.FrameData());
  ASSERT_EQ(val[3], 4);
  ASSERT_EQ(msg.NextFrame(), false);

  msg.FirstFrame();
  str = msg.FrameStr();
  ASSERT_STREQ(str, "abcdefg");
  msg.LastFrame();
  val = static_cast<int*>(msg.FrameData());
  ASSERT_EQ(val[2], 3);
}

TEST(MsgTest, AddFormatFrame) {
  int x = 5;
  Msg msg;
  msg.AddFormatFrame("i", 12);
  msg.AddFormatFrame("f", 10.f);
  msg.AddFormatFrame("s", "abc");
  msg.AddFormatFrame("p", &x);
  msg.AddFormatFrame("isfp", 12, "abc", 10.f, &x);

  msg.FirstFrame();
  int y;
  msg.ParseFormatFrame("i", &y);
  ASSERT_EQ(y, 12);
  ASSERT_EQ(msg.NextFrame(), true);

  float z;
  msg.ParseFormatFrame("f", &z);
  ASSERT_EQ(z, 10.f);
  ASSERT_EQ(msg.NextFrame(), true);

  char buf[10];
  msg.ParseFormatFrame("s", buf);
  ASSERT_STREQ(buf, "abc");
  ASSERT_EQ(msg.NextFrame(), true);

  int *p;
  msg.ParseFormatFrame("p", &p);
  ASSERT_EQ(p, &x);
  ASSERT_EQ(msg.NextFrame(), true);

  msg.ParseFormatFrame("isfp", &y, buf, &z, &p);
  ASSERT_EQ(y, 12);
  ASSERT_STREQ(buf, "abc");
  ASSERT_EQ(z, 10.f);
  ASSERT_EQ(p, &x);
}
<|endoftext|>


<commit_before>#include <sys/stat.h>
#include <float.h>

#include <iostream> // for standard output
#include <fstream> // for file input
#include <iomanip> // some nice printing functions
using namespace std;

#include <eigen3/Eigen/Dense> // linear algebra library
using namespace Eigen;

#include <boost/algorithm/string.hpp> // string manipulation library
#include <boost/program_options.hpp> // options parsing library
namespace po = boost::program_options;

#include "constants.h"
#include "qp-math.h"
#include "nv-math.h"
#include "printing.h"
#include "gates.h"

// random double from 0 to 1
inline double rnd(){ return std::rand()/(double)RAND_MAX; }

int main(int arg_num, const char *arg_vec[]) {

  // -----------------------------------------------------------------------------------------
  // Parse and process input options
  // -----------------------------------------------------------------------------------------

  // define parameters
  int cell_radius;
  double c13_abundance;
  string lattice_file = "lattice-cr[cell_radius]-s[seed].txt";

  int max_cluster_size;
  int ms;
  double Bz_in_gauss;

  bool perform_scan;
  int scan_bins;
  double f_DD;
  uint k_DD_int;
  double scan_time_in_ms;
  string larmor_file = "larmor-cr[cell_radius]-s[seed]-[ms].txt";
  string scan_file = "scan-cr[cell_radius]-s[seed]-[ms].txt";

  string output_dir;
  int seed;

  // define options
  po::options_description options("allowed options", 95);
  options.add_options()
    ("help,h", "produce help message")
    ("cell_radius", po::value<int>(&cell_radius)->default_value(6),
     "number of unit cells to simulate out from the NV center (along cell axes)")
    ("c13_abundance", po::value<double>(&c13_abundance)->default_value(0.0107,"0.0107"),
     "relative isotopic abundance of C-13")
    ("lattice_file", po::value<string>(&lattice_file),
     "specify file defining system configuration")

    ("max_cluster_size", po::value<int>(&max_cluster_size)->default_value(6),
     "maximum allowable size of C-13 clusters")
    ("ms", po::value<int>(&ms)->default_value(1),
     "NV center spin state used with |0> for an effective two-level system (+/-1)")
    ("Bz", po::value<double>(&Bz_in_gauss)->default_value(140.1,"140.1"),
     "strength of static magnetic field along the NV axis (in gauss)")

    ("scan", po::value<bool>(&perform_scan)->default_value(false),
     "perform coherence scan of effective larmor frequencies?")
    ("scan_bins", po::value<int>(&scan_bins)->default_value(200),
     "number of bins in coherence scanning range")
    ("f_DD", po::value<double>(&f_DD)->default_value(0.06,"0.06"),
     "magnitude of fourier component used in coherence scanning")
    ("k_DD", po::value<uint>(&k_DD_int)->default_value(1),
     "resonance harmonic used in coherence scanning (1 or 3)")
    ("scan_time", po::value<double>(&scan_time_in_ms)->default_value(1),
     "time for each coherence measurement (in milliseconds)")

    ("output_dir", po::value<string>(&output_dir)->default_value("./data"),
     "directory for storing data")
    ("seed", po::value<int>(&seed)->default_value(1),
     "seed for random number generator (>=1)")
    ;

  // collect inputs
  po::variables_map inputs;
  po::store(parse_command_line(arg_num, arg_vec, options), inputs);
  po::notify(inputs);

  // if requested, print help text
  if(inputs.count("help")){
    cout << options;
    return 0;
  }

  // print summary of inputs
  if(arg_num > 1){
    cout << "------------------------------------------------------------------" << endl;
    cout << "running " << arg_vec[0] << " with parameters:" << endl;
    for(int i = 1; i < arg_num; i++) {
      if(arg_vec[i][0] == '-') cout << endl;
      cout << arg_vec[i] << " ";
    }
    cout << endl;
    cout << "------------------------------------------------------------------" << endl;
    cout << endl;
  }

  bool set_cell_radius = !inputs["cell_radius"].defaulted();
  bool set_c13_abundance = !inputs["c13_abundance"].defaulted();
  bool input_lattice = inputs.count("lattice_file");

  // run a sanity check on inputs
  assert(cell_radius > 0);
  assert(c13_abundance >= 0 && c13_abundance <= 1);
  assert(!(input_lattice && set_cell_radius));
  assert(!(input_lattice && set_c13_abundance));

  assert(max_cluster_size > 0);
  assert(ms == 1 || ms == -1);
  assert(Bz_in_gauss >= 0);

  if(perform_scan){
    assert(scan_bins > 0);
    assert((f_DD > 0) && (f_DD < 1));
    assert((k_DD_int == 1) || (k_DD_int == 3));
    assert(scan_time_in_ms > 0);
  }

  assert(seed > 0);

  string lattice_path;
  if(input_lattice){
    lattice_path = lattice_file;
    // check that the given input file exists
    if(access(lattice_path.c_str(),F_OK) == -1){
      cout << "file does not exist: " << lattice_file << endl;
      return 1;
    }
  } else{ // if !input_lattice
    boost::replace_all(lattice_file, "[cell_radius]", to_string(cell_radius));
    boost::replace_all(lattice_file, "[seed]", to_string(seed));
    lattice_path = output_dir+"/"+lattice_file;
  }

  cout << "total number of lattice sites: " << int(pow(2*cell_radius,3)*cell_sites.size());
  cout << endl << endl;

  boost::replace_all(larmor_file, "[cell_radius]", to_string(cell_radius));
  boost::replace_all(larmor_file, "[seed]", to_string(seed));
  boost::replace_all(larmor_file, "[ms]", (ms > 0)?"up":"dn");
  string larmor_path = output_dir+"/"+larmor_file;

  boost::replace_all(scan_file, "[cell_radius]", to_string(cell_radius));
  boost::replace_all(scan_file, "[seed]", to_string(seed));
  boost::replace_all(scan_file, "[ms]", (ms > 0)?"up":"dn");
  string scan_path = output_dir+"/"+scan_file;

  // set variables based on iputs
  double Bz = Bz_in_gauss*gauss;
  harmonic k_DD;
  if(k_DD_int == 1) k_DD = first;
  if(k_DD_int == 3) k_DD = third;
  double scan_time = scan_time_in_ms*1e-3;

  srand(seed); // initialize random number generator
  mkdir(output_dir.c_str(),0777); // create data directory

  // -----------------------------------------------------------------------------------------
  // Construct spin lattice
  // -----------------------------------------------------------------------------------------

  // vector of C-13 nuclei
  vector<spin> nuclei;

  if(!input_lattice){

    // set positions of C-13 nuclei at lattice sites
    Vector3d cell_pos; // position of unit cell indexed by i,j,k
    for(int i = -cell_radius; i < cell_radius; i++){ // loop over all unit cell indices
      cell_pos(0) = i;
      for(int j = -cell_radius; j < cell_radius; j++){
        cell_pos(1) = j;
        for(int k = -cell_radius; k < cell_radius; k++){
          cell_pos(2) = k;
          for(uint ls = 0; ls < cell_sites.size(); ls++){ // loop over sites in unit cell
            if(rnd() <= c13_abundance){ // if we pass a check for C-13 isotopic abundance
              nuclei.push_back(spin(cell_pos+cell_sites.at(ls),gC13,s_vec));
            }
          }
        }
      }
    }

    // remove any C-13 atoms at the NV lattice sites
    nuclei.erase(remove(nuclei.begin(), nuclei.end(), spin(n.pos,gC13,s_vec)),
                 nuclei.end());
    nuclei.erase(remove(nuclei.begin(), nuclei.end(), spin(e(ms).pos,gC13,s_vec)),
                 nuclei.end());

    // write cell radius and C-13 positions to file
    ofstream lattice(lattice_path);
    lattice << "# cell_radius: " << cell_radius << endl;
    for(uint n = 0; n < nuclei.size(); n++){
      lattice << nuclei.at(n).pos(0) << ' '
             << nuclei.at(n).pos(1) << ' '
             << nuclei.at(n).pos(2) << endl;
    }
    lattice.close();

  } else { // if input_lattice

    string line;
    ifstream lattice(lattice_path);

    // get cell_radius
    getline(lattice,line,' ');
    getline(lattice,line,' ');
    getline(lattice,line);
    cell_radius = stoi(line);

    // get C-13 positions
    double x,y,z;
    while(getline(lattice,line,' ')){
      x = stod(line);
      getline(lattice,line,' ');
      y = stod(line);
      getline(lattice,line);
      z = stod(line);
      nuclei.push_back(spin((Vector3d() << x,y,z).finished(),gC13,s_vec));
    }
    lattice.close();

  }

  // assert that no C-13 spins lie at the NV lattice sites
  if(in_vector(spin(n.pos,gC13,s_vec),nuclei)
     || in_vector(spin(e(ms).pos,gC13,s_vec),nuclei)){
    cout << "we have a C-13 nucleus at one of the NV lattice sites!" << endl;
    return 2;
  }

  // -----------------------------------------------------------------------------------------
  // Cluster C-13 nuclei
  // -----------------------------------------------------------------------------------------

  double cluster_coupling_guess = 100;
  double dcc_cutoff = 1e-5;

  // get cluster_coupling for which the largest cluster size is >= max_cluster_size
  double cluster_coupling = find_target_coupling(nuclei,max_cluster_size,
                                                 cluster_coupling_guess,dcc_cutoff);
  vector<vector<spin>> clusters = get_clusters(nuclei,cluster_coupling);

  // if largest cluster size > max_cluster_size,
  //   which can occur because it is impossible to exactly match max_cluster_size,
  //   find largest cluster_coupling for which the largest cluster size is < max_cluster_size
  if(largest_cluster_size(clusters) > max_cluster_size){
    int cluster_size_target
      = largest_cluster_size(get_clusters(nuclei,cluster_coupling+dcc_cutoff));
    cluster_coupling = find_target_coupling(nuclei,cluster_size_target,
                                            cluster_coupling,dcc_cutoff);
    clusters = get_clusters(nuclei,cluster_coupling);
  }

  cout << nuclei.size() << " nuclei grouped into " << clusters.size() << " clusters" << endl;
  cout << "largest cluster size: " << largest_cluster_size(clusters) << endl;
  cout << "cluster coupling factor: " << cluster_coupling << " Hz" << endl;
  cout << endl;

  // -----------------------------------------------------------------------------------------
  // Coherence scan
  // -----------------------------------------------------------------------------------------

  if(perform_scan){

    // idenfy effictive larmor frequencies and NV coupling strengths
    VectorXd w_larmor = VectorXd::Zero(nuclei.size());
    VectorXd A_perp = VectorXd::Zero(nuclei.size());

    double w_max = 0, w_min = DBL_MAX;
    for(uint n = 0; n < nuclei.size(); n++){
      Vector3d A_n = A(nuclei.at(n),ms);
      w_larmor(n) = effective_larmor(nuclei.at(n),Bz*zhat,A_n,ms);
      A_perp(n) = (A_n-dot(A_n,zhat)*zhat).norm();

      if(w_larmor(n) < w_min) w_min = w_larmor(n);
      if(w_larmor(n) > w_max) w_max = w_larmor(n);
    }

    ofstream larmor(larmor_path);
    larmor << "# f_DD: " << f_DD << endl;
    larmor << "# k_DD: " << k_DD_int << endl;
    larmor << "# w A_perp" << endl;
    for(uint n = 0; n < nuclei.size(); n++){
      larmor << w_larmor(n) << " " << A_perp(n) << endl;
    }
    larmor.close();

    // perform coherence scan
    cout << "beginning coherence scan" << endl;
    VectorXd w_scan = VectorXd::Zero(scan_bins);
    VectorXd coherence = VectorXd::Zero(scan_bins);

    double w_range = w_max - w_min;
    double w_start = w_min - w_range/10;
    double w_end = w_max + w_range/10;
    for(uint i = 0; i < scan_bins; i++){
      w_scan(i) = w_start + i*(w_end-w_start)/scan_bins;
      coherence(i) =
        coherence_measurement(ms, clusters, w_scan(i), k_DD, f_DD, Bz, scan_time);
      cout << "(" << i+1 << "/" << scan_bins << ") "
           << w_scan(i) << " " << coherence(i) << endl;
    }

    ofstream scan(scan_path);
    scan << "# f_DD: " << f_DD << endl;
    scan << "# k_DD: " << k_DD_int << endl;
    scan << "# w L" << endl;
    for(uint i = 0; i < scan_bins; i++){
      scan << w_scan(i) << " " << coherence(i) << endl;
    }
    scan.close();

  }

}

<commit_msg>file naming changes<commit_after>#include <sys/stat.h>
#include <float.h>

#include <iostream> // for standard output
#include <fstream> // for file input
#include <iomanip> // some nice printing functions
using namespace std;

#include <eigen3/Eigen/Dense> // linear algebra library
using namespace Eigen;

#include <boost/algorithm/string.hpp> // string manipulation library
#include <boost/program_options.hpp> // options parsing library
namespace po = boost::program_options;

#include "constants.h"
#include "qp-math.h"
#include "nv-math.h"
#include "printing.h"
#include "gates.h"

// random double from 0 to 1
inline double rnd(){ return std::rand()/(double)RAND_MAX; }

int main(int arg_num, const char *arg_vec[]) {

  // -----------------------------------------------------------------------------------------
  // Parse and process input options
  // -----------------------------------------------------------------------------------------

  // define parameters
  int cell_radius;
  double c13_abundance;
  string lattice_file = "lattice-cr[cell_radius]-s[seed].txt";

  int max_cluster_size;
  int ms;
  double Bz_in_gauss;

  bool perform_scan;
  int scan_bins;
  double f_DD;
  uint k_DD_int;
  double scan_time_in_us;
  string file_suffix = "-[ms]-cs[cluster_size]-cr[cell_radius]-s[seed].txt";

  string output_dir;
  int seed;

  // define options
  po::options_description options("Allowed options", 95);
  options.add_options()
    ("help,h", "produce help message")
    ("cell_radius", po::value<int>(&cell_radius)->default_value(7),
     "number of unit cells to simulate out from the NV center (along cell axes)")
    ("c13_abundance", po::value<double>(&c13_abundance)->default_value(0.0107,"0.0107"),
     "relative isotopic abundance of C-13")
    ("lattice_file", po::value<string>(&lattice_file),
     "specify file defining system configuration")

    ("max_cluster_size", po::value<int>(&max_cluster_size)->default_value(7),
     "maximum allowable size of C-13 clusters")
    ("ms", po::value<int>(&ms)->default_value(1),
     "NV center spin state used with |0> for an effective two-level system (+/-1)")
    ("Bz", po::value<double>(&Bz_in_gauss)->default_value(140.1,"140.1"),
     "strength of static magnetic field along the NV axis (in gauss)")

    ("scan", po::value<bool>(&perform_scan)->default_value(false),
     "perform coherence scan of effective larmor frequencies?")
    ("scan_bins", po::value<int>(&scan_bins)->default_value(200),
     "number of bins in coherence scanning range")
    ("f_DD", po::value<double>(&f_DD)->default_value(0.06,"0.06"),
     "magnitude of fourier component used in coherence scanning")
    ("k_DD", po::value<uint>(&k_DD_int)->default_value(1),
     "resonance harmonic used in coherence scanning (1 or 3)")
    ("scan_time", po::value<double>(&scan_time_in_us)->default_value(100),
     "time for each coherence measurement (in microseconds)")

    ("output_dir", po::value<string>(&output_dir)->default_value("./data"),
     "directory for storing data")
    ("seed", po::value<int>(&seed)->default_value(1),
     "seed for random number generator (>=1)")
    ;

  // collect inputs
  po::variables_map inputs;
  po::store(parse_command_line(arg_num, arg_vec, options), inputs);
  po::notify(inputs);

  // if requested, print help text
  if(inputs.count("help")){
    cout << options;
    return 0;
  }

  // print summary of inputs
  if(arg_num > 1){
    cout << "------------------------------------------------------------------" << endl;
    cout << "running " << arg_vec[0] << " with parameters:" << endl;
    for(int i = 1; i < arg_num; i++) {
      if(arg_vec[i][0] == '-') cout << endl;
      cout << arg_vec[i] << " ";
    }
    cout << endl;
    cout << "------------------------------------------------------------------" << endl;
    cout << endl;
  }

  bool set_cell_radius = !inputs["cell_radius"].defaulted();
  bool set_c13_abundance = !inputs["c13_abundance"].defaulted();
  bool input_lattice = inputs.count("lattice_file");

  // run a sanity check on inputs
  assert(cell_radius > 0);
  assert(c13_abundance >= 0 && c13_abundance <= 1);
  assert(!(input_lattice && set_cell_radius));
  assert(!(input_lattice && set_c13_abundance));

  assert(max_cluster_size > 0);
  assert(ms == 1 || ms == -1);
  assert(Bz_in_gauss >= 0);

  if(perform_scan){
    assert(scan_bins > 0);
    assert((f_DD > 0) && (f_DD < 1));
    assert((k_DD_int == 1) || (k_DD_int == 3));
    assert(scan_time_in_us > 0);
  }

  assert(seed > 0);

  string lattice_path;
  if(input_lattice){
    lattice_path = lattice_file;
    // check that the given input file exists
    if(access(lattice_path.c_str(),F_OK) == -1){
      cout << "file does not exist: " << lattice_file << endl;
      return 1;
    }
  } else{ // if !input_lattice
    boost::replace_all(lattice_file, "[cell_radius]", to_string(cell_radius));
    boost::replace_all(lattice_file, "[seed]", to_string(seed));
    lattice_path = output_dir+"/"+lattice_file;
  }

  cout << "Total number of lattice sites: " << int(pow(2*cell_radius,3)*cell_sites.size());
  cout << endl << endl;

  // set variables based on iputs
  double Bz = Bz_in_gauss*gauss;
  harmonic k_DD;
  if(k_DD_int == 1) k_DD = first;
  if(k_DD_int == 3) k_DD = third;
  double scan_time = scan_time_in_us*1e-6;

  srand(seed); // initialize random number generator
  mkdir(output_dir.c_str(),0777); // create data directory

  // -----------------------------------------------------------------------------------------
  // Construct spin lattice
  // -----------------------------------------------------------------------------------------

  // vector of C-13 nuclei
  vector<spin> nuclei;

  if(!input_lattice){

    // set positions of C-13 nuclei at lattice sites
    Vector3d cell_pos; // position of unit cell indexed by i,j,k
    for(int i = -cell_radius; i < cell_radius; i++){ // loop over all unit cell indices
      cell_pos(0) = i;
      for(int j = -cell_radius; j < cell_radius; j++){
        cell_pos(1) = j;
        for(int k = -cell_radius; k < cell_radius; k++){
          cell_pos(2) = k;
          for(uint ls = 0; ls < cell_sites.size(); ls++){ // loop over sites in unit cell
            if(rnd() <= c13_abundance){ // if we pass a check for C-13 isotopic abundance
              nuclei.push_back(spin(cell_pos+cell_sites.at(ls),gC13,s_vec));
            }
          }
        }
      }
    }

    // remove any C-13 atoms at the NV lattice sites
    nuclei.erase(remove(nuclei.begin(), nuclei.end(), spin(n.pos,gC13,s_vec)),
                 nuclei.end());
    nuclei.erase(remove(nuclei.begin(), nuclei.end(), spin(e(ms).pos,gC13,s_vec)),
                 nuclei.end());

    // write cell radius and C-13 positions to file
    ofstream lattice(lattice_path);
    lattice << "# cell_radius: " << cell_radius << endl;
    for(uint n = 0; n < nuclei.size(); n++){
      lattice << nuclei.at(n).pos(0) << ' '
             << nuclei.at(n).pos(1) << ' '
             << nuclei.at(n).pos(2) << endl;
    }
    lattice.close();

  } else { // if input_lattice

    string line;
    ifstream lattice(lattice_path);

    // get cell_radius
    getline(lattice,line,' ');
    getline(lattice,line,' ');
    getline(lattice,line);
    cell_radius = stoi(line);

    // get C-13 positions
    double x,y,z;
    while(getline(lattice,line,' ')){
      x = stod(line);
      getline(lattice,line,' ');
      y = stod(line);
      getline(lattice,line);
      z = stod(line);
      nuclei.push_back(spin((Vector3d() << x,y,z).finished(),gC13,s_vec));
    }
    lattice.close();

  }

  // assert that no C-13 spins lie at the NV lattice sites
  if(in_vector(spin(n.pos,gC13,s_vec),nuclei)
     || in_vector(spin(e(ms).pos,gC13,s_vec),nuclei)){
    cout << "we have a C-13 nucleus at one of the NV lattice sites!" << endl;
    return 2;
  }

  // -----------------------------------------------------------------------------------------
  // Cluster C-13 nuclei
  // -----------------------------------------------------------------------------------------

  cout << "Clustering " << nuclei.size() << " nuclei" << endl;
  double cluster_coupling_guess = 100;
  double dcc_cutoff = 1e-5;

  // get cluster_coupling for which the largest cluster size is >= max_cluster_size
  double cluster_coupling = find_target_coupling(nuclei,max_cluster_size,
                                                 cluster_coupling_guess,dcc_cutoff);
  vector<vector<spin>> clusters = get_clusters(nuclei,cluster_coupling);

  // if largest cluster size > max_cluster_size,
  //   which can occur because it is impossible to exactly match max_cluster_size,
  //   find largest cluster_coupling for which the largest cluster size is < max_cluster_size
  if(largest_cluster_size(clusters) > max_cluster_size){
    int cluster_size_target
      = largest_cluster_size(get_clusters(nuclei,cluster_coupling+dcc_cutoff));
    cluster_coupling = find_target_coupling(nuclei,cluster_size_target,
                                            cluster_coupling,dcc_cutoff);
    clusters = get_clusters(nuclei,cluster_coupling);
  }

  max_cluster_size = largest_cluster_size(clusters);
  cout << "Nuclei grouped into " << clusters.size() << " clusters" << endl;
  cout << "Maximum cluster size: " << max_cluster_size << endl;
  cout << "Cluster coupling factor: " << cluster_coupling << " Hz" << endl;
  cout << endl;

  // -----------------------------------------------------------------------------------------
  // Coherence scan
  // -----------------------------------------------------------------------------------------

  if(perform_scan){

    boost::replace_all(file_suffix, "[ms]", (ms > 0)?"up":"dn");
    boost::replace_all(file_suffix, "[cluster_size]", to_string(max_cluster_size));
    boost::replace_all(file_suffix, "[cell_radius]", to_string(cell_radius));
    boost::replace_all(file_suffix, "[seed]", to_string(seed));

    string larmor_path = output_dir+"/larmor"+file_suffix;
    string scan_path = output_dir+"/scan"+file_suffix;

    stringstream file_header;
    file_header << "# maxium cluster size: " << max_cluster_size << endl;
    file_header << "# cluster coupling factor (Hz): " << cluster_coupling << endl;
    file_header << "# f_DD: " << f_DD << endl;
    file_header << "# k_DD: " << k_DD_int << endl;
    file_header << "# scan time: " << scan_time << endl;

    // idenfy effictive larmor frequencies and NV coupling strengths
    VectorXd w_larmor = VectorXd::Zero(nuclei.size());
    VectorXd A_perp = VectorXd::Zero(nuclei.size());

    double w_max = 0, w_min = DBL_MAX;
    for(uint n = 0; n < nuclei.size(); n++){
      Vector3d A_n = A(nuclei.at(n),ms);
      w_larmor(n) = effective_larmor(nuclei.at(n),Bz*zhat,A_n,ms);
      A_perp(n) = (A_n-dot(A_n,zhat)*zhat).norm();

      if(w_larmor(n) < w_min) w_min = w_larmor(n);
      if(w_larmor(n) > w_max) w_max = w_larmor(n);
    }

    ofstream larmor(larmor_path);
    larmor << file_header.str();
    larmor << endl << "# w_larmor A_perp" << endl;
    for(uint n = 0; n < nuclei.size(); n++){
      larmor << w_larmor(n) << " " << A_perp(n) << endl;
    }
    larmor.close();

    // perform coherence scan
    cout << "Beginning coherence scan" << endl;
    VectorXd w_scan = VectorXd::Zero(scan_bins);
    VectorXd coherence = VectorXd::Zero(scan_bins);

    double w_range = w_max - w_min;
    double w_start = w_min - w_range/10;
    double w_end = w_max + w_range/10;
    for(int i = 0; i < scan_bins; i++){
      w_scan(i) = w_start + i*(w_end-w_start)/scan_bins;
      coherence(i) =
        coherence_measurement(ms, clusters, w_scan(i), k_DD, f_DD, Bz, scan_time);
      cout << "(" << i+1 << "/" << scan_bins << ") "
           << w_scan(i) << " " << coherence(i) << endl;
    }

    ofstream scan(scan_path);
    scan << file_header.str();
    scan << endl << "# w_scan coherence" << endl;
    for(int i = 0; i < scan_bins; i++){
      scan << int(w_scan(i)) << " " << coherence(i) << endl;
    }
    scan.close();

  }

}

<|endoftext|>


<commit_before>#include "generated.hpp"
#include "impl_test_interface.hpp"
#include "checkpoint.hpp"

namespace warpcoil
{
    namespace
    {
        template <class Handler, class TotalResult>
        struct wait_for_all_state
        {
            Handler handler;
            TotalResult total_result;
            std::size_t completed = 0;
            bool errored = false;

            template <class CompletionToken>
            wait_for_all_state(CompletionToken &&token)
                : handler(std::forward<CompletionToken>(token))
            {
            }
        };

        template <std::size_t Index, std::size_t MethodCount, class Method, class State>
        void start_method(Method &&method, std::shared_ptr<State> const &state)
        {
            method.start([state](boost::system::error_code ec, typename std::decay<Method>::type::result element)
                         {
                             if (!ec)
                             {
                                 std::get<Index>(state->total_result) = std::move(element);
                                 ++state->completed;
                                 if (state->completed == MethodCount)
                                 {
                                     state->handler(ec, std::move(state->total_result));
                                 }
                             }
                             else if (!Si::exchange(state->errored, true))
                             {
                                 state->handler(ec, {});
                             }
                         });
        }

        template <class State, class... Methods, std::size_t... Indices>
        void start_methods(std::shared_ptr<State> state, std::integer_sequence<std::size_t, Indices...>,
                           Methods &&... methods)
        {
            Si::unit dummy[] = {(start_method<Indices, sizeof...(Methods)>(methods, state), Si::unit())...};
            Si::ignore_unused_variable_warning(dummy);
        }

        template <class... Methods, class CompletionToken>
        auto wait_for_all(CompletionToken &&token, Methods &&... methods)
        {
            BOOST_STATIC_ASSERT(sizeof...(Methods) >= 1);
            using total_result = std::tuple<typename std::decay<Methods>::type::result...>;
            using handler_type = typename boost::asio::handler_type<decltype(token), void(boost::system::error_code,
                                                                                total_result)>::type;
            auto state =
                std::make_shared<wait_for_all_state<handler_type, total_result>>(std::forward<CompletionToken>(token));
            boost::asio::async_result<handler_type> result(state->handler);
            start_methods(std::move(state), std::make_integer_sequence<std::size_t, sizeof...(Methods)>(),
                          std::forward<Methods>(methods)...);
            return result.get();
        }

        template <class Callback>
        struct result_of_callback;

        template <class Result>
        struct result_of_callback<std::function<void(boost::system::error_code, Result)>>
        {
            using type = Result;
        };

        template <class Result, class Start>
        struct method_holder
        {
            using result = Result;
            Start start;
        };

        template <typename F, typename Tuple, size_t... S>
        decltype(auto) apply_tuple_impl(F &&fn, Tuple &&t, std::index_sequence<S...>)
        {
            return std::forward<F>(fn)(std::get<S>(std::forward<Tuple>(t))...);
        }

        template <typename F, typename Tuple>
        decltype(auto) apply_from_tuple(F &&fn, Tuple &&t)
        {
            std::size_t constexpr tSize = std::tuple_size<typename std::remove_reference<Tuple>::type>::value;
            return apply_tuple_impl(std::forward<F>(fn), std::forward<Tuple>(t), std::make_index_sequence<tSize>());
        }

        template <class Interface, class Parameter0, class... Parameters, class... Arguments>
        auto method(Interface &callee, void (Interface::*method_ptr)(Parameter0, Parameters...),
                    Arguments &&... arguments)
        {
            auto start = [&callee, method_ptr, arguments = std::make_tuple(std::forward<Arguments>(arguments)...) ](
                auto &&callback) mutable
            {
                apply_from_tuple(
                    [&](auto &&... arguments)
                    {
                        (callee.*method_ptr)(std::forward<decltype(arguments)>(arguments)...,
                                             std::forward<decltype(callback)>(callback));
                    },
                    std::move(arguments));
            };
            using std_function_callback = typename boost::mpl::vector<Parameters...>::back;
            using result = typename result_of_callback<std_function_callback>::type;
            return method_holder<result, decltype(start)>{std::move(start)};
        }
    }
}

BOOST_AUTO_TEST_CASE(wait_for_all_1)
{
    warpcoil::impl_test_interface server;
    warpcoil::checkpoint got_result;
    got_result.enable();
    warpcoil::wait_for_all(
        [&got_result](boost::system::error_code ec, std::tuple<std::string> result)
        {
            got_result.enter();
            BOOST_CHECK(!ec);
            BOOST_CHECK_EQUAL("Hello123", std::get<0>(result));
        },
        warpcoil::method(server, &warpcoil::impl_test_interface::utf8, "Hello"));
}

BOOST_AUTO_TEST_CASE(wait_for_all_2)
{
    warpcoil::impl_test_interface server;
    warpcoil::checkpoint got_result;
    got_result.enable();
    warpcoil::wait_for_all(
        [&got_result](boost::system::error_code ec, std::tuple<std::string, std::uint16_t> result)
        {
            got_result.enter();
            BOOST_CHECK(!ec);
            BOOST_CHECK_EQUAL("Hello123", std::get<0>(result));
            BOOST_CHECK_EQUAL(123u, std::get<1>(result));
        },
        warpcoil::method(server, &warpcoil::impl_test_interface::utf8, "Hello"),
        warpcoil::method(server, &warpcoil::impl_test_interface::atypical_int, 123));
}
<commit_msg>use boost::mpl::back correctly<commit_after>#include "generated.hpp"
#include "impl_test_interface.hpp"
#include "checkpoint.hpp"
#include <boost/mpl/back.hpp>

namespace warpcoil
{
    namespace
    {
        template <class Handler, class TotalResult>
        struct wait_for_all_state
        {
            Handler handler;
            TotalResult total_result;
            std::size_t completed = 0;
            bool errored = false;

            template <class CompletionToken>
            wait_for_all_state(CompletionToken &&token)
                : handler(std::forward<CompletionToken>(token))
            {
            }
        };

        template <std::size_t Index, std::size_t MethodCount, class Method, class State>
        void start_method(Method &&method, std::shared_ptr<State> const &state)
        {
            method.start([state](boost::system::error_code ec, typename std::decay<Method>::type::result element)
                         {
                             if (!ec)
                             {
                                 std::get<Index>(state->total_result) = std::move(element);
                                 ++state->completed;
                                 if (state->completed == MethodCount)
                                 {
                                     state->handler(ec, std::move(state->total_result));
                                 }
                             }
                             else if (!Si::exchange(state->errored, true))
                             {
                                 state->handler(ec, {});
                             }
                         });
        }

        template <class State, class... Methods, std::size_t... Indices>
        void start_methods(std::shared_ptr<State> state, std::integer_sequence<std::size_t, Indices...>,
                           Methods &&... methods)
        {
            Si::unit dummy[] = {(start_method<Indices, sizeof...(Methods)>(methods, state), Si::unit())...};
            Si::ignore_unused_variable_warning(dummy);
        }

        template <class... Methods, class CompletionToken>
        auto wait_for_all(CompletionToken &&token, Methods &&... methods)
        {
            BOOST_STATIC_ASSERT(sizeof...(Methods) >= 1);
            using total_result = std::tuple<typename std::decay<Methods>::type::result...>;
            using handler_type = typename boost::asio::handler_type<decltype(token), void(boost::system::error_code,
                                                                                total_result)>::type;
            auto state =
                std::make_shared<wait_for_all_state<handler_type, total_result>>(std::forward<CompletionToken>(token));
            boost::asio::async_result<handler_type> result(state->handler);
            start_methods(std::move(state), std::make_integer_sequence<std::size_t, sizeof...(Methods)>(),
                          std::forward<Methods>(methods)...);
            return result.get();
        }

        template <class Callback>
        struct result_of_callback;

        template <class Result>
        struct result_of_callback<std::function<void(boost::system::error_code, Result)>>
        {
            using type = Result;
        };

        template <class Result, class Start>
        struct method_holder
        {
            using result = Result;
            Start start;
        };

        template <typename F, typename Tuple, size_t... S>
        decltype(auto) apply_tuple_impl(F &&fn, Tuple &&t, std::index_sequence<S...>)
        {
            return std::forward<F>(fn)(std::get<S>(std::forward<Tuple>(t))...);
        }

        template <typename F, typename Tuple>
        decltype(auto) apply_from_tuple(F &&fn, Tuple &&t)
        {
            std::size_t constexpr tSize = std::tuple_size<typename std::remove_reference<Tuple>::type>::value;
            return apply_tuple_impl(std::forward<F>(fn), std::forward<Tuple>(t), std::make_index_sequence<tSize>());
        }

        template <class Interface, class Parameter0, class... Parameters, class... Arguments>
        auto method(Interface &callee, void (Interface::*method_ptr)(Parameter0, Parameters...),
                    Arguments &&... arguments)
        {
            auto start = [&callee, method_ptr, arguments = std::make_tuple(std::forward<Arguments>(arguments)...) ](
                auto &&callback) mutable
            {
                apply_from_tuple(
                    [&](auto &&... arguments)
                    {
                        (callee.*method_ptr)(std::forward<decltype(arguments)>(arguments)...,
                                             std::forward<decltype(callback)>(callback));
                    },
                    std::move(arguments));
            };
            using std_function_callback = typename boost::mpl::back<boost::mpl::vector<Parameters...>>::type;
            using result = typename result_of_callback<std_function_callback>::type;
            return method_holder<result, decltype(start)>{std::move(start)};
        }
    }
}

BOOST_AUTO_TEST_CASE(wait_for_all_1)
{
    warpcoil::impl_test_interface server;
    warpcoil::checkpoint got_result;
    got_result.enable();
    warpcoil::wait_for_all(
        [&got_result](boost::system::error_code ec, std::tuple<std::string> result)
        {
            got_result.enter();
            BOOST_CHECK(!ec);
            BOOST_CHECK_EQUAL("Hello123", std::get<0>(result));
        },
        warpcoil::method(server, &warpcoil::impl_test_interface::utf8, "Hello"));
}

BOOST_AUTO_TEST_CASE(wait_for_all_2)
{
    warpcoil::impl_test_interface server;
    warpcoil::checkpoint got_result;
    got_result.enable();
    warpcoil::wait_for_all(
        [&got_result](boost::system::error_code ec, std::tuple<std::string, std::uint16_t> result)
        {
            got_result.enter();
            BOOST_CHECK(!ec);
            BOOST_CHECK_EQUAL("Hello123", std::get<0>(result));
            BOOST_CHECK_EQUAL(123u, std::get<1>(result));
        },
        warpcoil::method(server, &warpcoil::impl_test_interface::utf8, "Hello"),
        warpcoil::method(server, &warpcoil::impl_test_interface::atypical_int, static_cast<boost::uint16_t>(123)));
}
<|endoftext|>


<commit_before>#ifndef AMANZI_TEST_PK_BC_FACTORY_HH_
#define AMANZI_TEST_PK_BC_FACTORY_HH_

/* -------------------------------------------------------------------------
ATS

Author: ...
    Ethan Coon (ATS version) (ecoon@lanl.gov)

*/

#include "Teuchos_RCP.hpp"
#include "Teuchos_ParameterList.hpp"

#include "bc_factory.hh"

namespace Amanzi {
namespace TestPKs {

class TestPKBCFactory : public Amanzi::BCFactory {

public:
  TestPKBCFactory(const Teuchos::RCP<const AmanziMesh::Mesh> &mesh,
                  const Teuchos::ParameterList& plist)
      : Amanzi::BCFactory(mesh,plist) {}

  Teuchos::RCP<Functions::BoundaryFunction> CreateDirichlet() const {
    return CreateWithFunction("dirichlet", "boundary data");
  }

  Teuchos::RCP<Functions::BoundaryFunction> CreateNeumann() const {
    return CreateWithFunction("neumann", "outward flux");
  }

};

}  // namespace
}  // namespace

#endif // AMANZI_FLOW_BC_FACTORY_HH_
<commit_msg>bug fix in test bc factory<commit_after>#ifndef AMANZI_TEST_PK_BC_FACTORY_HH_
#define AMANZI_TEST_PK_BC_FACTORY_HH_

/* -------------------------------------------------------------------------
ATS

Author: ...
    Ethan Coon (ATS version) (ecoon@lanl.gov)

*/

#include "Teuchos_RCP.hpp"
#include "Teuchos_ParameterList.hpp"

#include "bc_factory.hh"

namespace Amanzi {
namespace TestPKs {

class TestPKBCFactory : public BCFactory {

public:
  TestPKBCFactory(const Teuchos::RCP<const AmanziMesh::Mesh> &mesh,
                  const Teuchos::ParameterList& plist) :
	BCFactory(mesh,plist) {}

  Teuchos::RCP<Functions::BoundaryFunction> CreateDirichlet() const {
    return CreateWithFunction("dirichlet", "boundary data");
  }

  Teuchos::RCP<Functions::BoundaryFunction> CreateNeumann() const {
    return CreateWithFunction("neumann", "outward flux");
  }

};

}  // namespace
}  // namespace

#endif // AMANZI_FLOW_BC_FACTORY_HH_
<|endoftext|>


<commit_before>/*
  WiFi.cpp - Library for Arduino Wifi shield.
  Copyright (c) 2011-2014 Arduino.  All right reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.

  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.

  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include "ESP8266WiFi.h"
extern "C" {
#include "c_types.h"
#include "ets_sys.h"
#include "os_type.h"
#include "osapi.h"
#include "mem.h"
#include "user_interface.h"
#include "lwip/opt.h"
#include "lwip/err.h"
#include "lwip/dns.h"
}

extern "C" void esp_schedule();
extern "C" void esp_yield();

ESP8266WiFiClass::ESP8266WiFiClass()
{

}

void ESP8266WiFiClass::mode(WiFiMode m)
{
    ETS_UART_INTR_DISABLE();
    wifi_set_opmode(m);
    ETS_UART_INTR_ENABLE();
}

int ESP8266WiFiClass::begin(const char* ssid)
{
    return begin(ssid, 0);
}

int ESP8266WiFiClass::begin(const char* ssid, const char *passphrase)
{
    if (wifi_get_opmode() == WIFI_AP)
    {
        // turn on AP+STA mode
        mode(WIFI_AP_STA);
    }

    struct station_config conf;
    strcpy(reinterpret_cast<char*>(conf.ssid), ssid);
    if (passphrase)
        strcpy(reinterpret_cast<char*>(conf.password), passphrase);
    else
        *conf.password = 0;

    ETS_UART_INTR_DISABLE();
    wifi_station_set_config(&conf);
    wifi_station_connect();
    ETS_UART_INTR_ENABLE();
    wifi_station_dhcp_start();
    return 0;
}

void ESP8266WiFiClass::config(IPAddress local_ip, IPAddress gateway, IPAddress subnet)
{
    struct ip_info info;
    info.ip.addr = static_cast<uint32_t>(local_ip);
    info.gw.addr = static_cast<uint32_t>(gateway);
    info.netmask.addr = static_cast<uint32_t>(subnet);

    wifi_station_dhcpc_stop();
    wifi_set_ip_info(STATION_IF, &info);
}

int ESP8266WiFiClass::disconnect()
{
    struct station_config conf;
    *conf.ssid = 0;
    *conf.password = 0;
    ETS_UART_INTR_DISABLE();
    wifi_station_set_config(&conf);
    wifi_station_disconnect();
    ETS_UART_INTR_ENABLE();
    return 0;
}

void ESP8266WiFiClass::softAP(const char* ssid)
{
    softAP(ssid, 0);
}

void ESP8266WiFiClass::softAP(const char* ssid, const char* passphrase)
{
    if (wifi_get_opmode() == WIFI_STA)
    {
        // turn on AP+STA mode
        mode(WIFI_AP_STA);
    }

    struct softap_config conf;
    wifi_softap_get_config(&conf);
    strcpy(reinterpret_cast<char*>(conf.ssid), ssid);
    conf.channel = 1;

    if (!passphrase || strlen(passphrase) == 0)
    {
        conf.authmode = AUTH_OPEN;
        *conf.password = 0;
    }
    else
    {
        conf.authmode = AUTH_WPA2_PSK;
        strcpy(reinterpret_cast<char*>(conf.password), passphrase);
    }

    ETS_UART_INTR_DISABLE();
    wifi_softap_set_config(&conf);
    ETS_UART_INTR_ENABLE();
}

uint8_t* ESP8266WiFiClass::macAddress(uint8_t* mac)
{
    wifi_get_macaddr(STATION_IF, mac);
    return mac;
}

uint8_t* softAPmacAddress(uint8_t* mac)
{
    wifi_get_macaddr(SOFTAP_IF, mac);
    return mac;
}

IPAddress ESP8266WiFiClass::localIP()
{
    struct ip_info ip;
    wifi_get_ip_info(STATION_IF, &ip);
    return IPAddress(ip.ip.addr);
}

IPAddress softAPIP()
{
    struct ip_info ip;
    wifi_get_ip_info(SOFTAP_IF, &ip);
    return IPAddress(ip.ip.addr);
}

IPAddress ESP8266WiFiClass::subnetMask()
{
    struct ip_info ip;
    wifi_get_ip_info(STATION_IF, &ip);
    return IPAddress(ip.netmask.addr);
}

IPAddress ESP8266WiFiClass::gatewayIP()
{
    struct ip_info ip;
    wifi_get_ip_info(STATION_IF, &ip);
    return IPAddress(ip.gw.addr);
}

char* ESP8266WiFiClass::SSID()
{
    static struct station_config conf;
    wifi_station_get_config(&conf);
    return reinterpret_cast<char*>(conf.ssid);
}

// uint8_t* ESP8266WiFiClass::BSSID(uint8_t* bssid)
// {
//  uint8_t* _bssid = WiFiDrv::getCurrentBSSID();
//  memcpy(bssid, _bssid, WL_MAC_ADDR_LENGTH);
//     return bssid;
// }

// int32_t ESP8266WiFiClass::RSSI()
// {
//     return WiFiDrv::getCurrentRSSI();
// }

// uint8_t ESP8266WiFiClass::encryptionType()
// {
//     return WiFiDrv::getCurrentEncryptionType();
// }

void ESP8266WiFiClass::_scanDone(void* result, int status)
{
    if (status != OK)
    {
        ESP8266WiFiClass::_scanCount = 0;
        ESP8266WiFiClass::_scanResult = 0;
    }
    else
    {
      ESP8266WiFiClass::_scanResult = result;
      int i = 0;
      for (bss_info* it = reinterpret_cast<bss_info*>(result); it; it = STAILQ_NEXT(it, next), ++i);
      ESP8266WiFiClass::_scanCount = i;
    }
    esp_schedule();
}

int8_t ESP8266WiFiClass::scanNetworks()
{
    if (wifi_get_opmode() == WIFI_AP)
    {
        mode(WIFI_AP_STA);
    }
    int status = wifi_station_get_connect_status();
    if (status != STATION_GOT_IP && status != STATION_IDLE)
    {
        disconnect();
    }

    struct scan_config config;
    config.ssid = 0;
    config.bssid = 0;
    config.channel = 0;
    wifi_station_scan(&config, reinterpret_cast<scan_done_cb_t>(&ESP8266WiFiClass::_scanDone));
    esp_yield();
    return ESP8266WiFiClass::_scanCount;
}

void * ESP8266WiFiClass::_getScanInfoByIndex(int i)
{
    if (!ESP8266WiFiClass::_scanResult || i > ESP8266WiFiClass::_scanCount)
        return 0;

    struct bss_info* it = reinterpret_cast<struct bss_info*>(ESP8266WiFiClass::_scanResult);
    for (; i; it = STAILQ_NEXT(it, next), --i);
    if (!it)
        return 0;

    return it;
}

char* ESP8266WiFiClass::SSID(uint8_t i)
{
    struct bss_info* it = reinterpret_cast<struct bss_info*>(_getScanInfoByIndex(i));
    if (!it)
        return 0;

    return reinterpret_cast<char*>(it->ssid);
}

int32_t ESP8266WiFiClass::RSSI(uint8_t i)
{
    struct bss_info* it = reinterpret_cast<struct bss_info*>(_getScanInfoByIndex(i));
    if (!it)
        return 0;

    return it->rssi;
}

uint8_t ESP8266WiFiClass::encryptionType(uint8_t i)
{
    struct bss_info* it = reinterpret_cast<struct bss_info*>(_getScanInfoByIndex(i));
    if (!it)
        return -1;

    int authmode = it->authmode;
    if (authmode == AUTH_OPEN)
        return ENC_TYPE_NONE;
    if (authmode == AUTH_WEP)
        return ENC_TYPE_WEP;
    if (authmode == AUTH_WPA_PSK || authmode == AUTH_WPA_WPA2_PSK)  // fixme: is this correct?
        return ENC_TYPE_TKIP;
    if (authmode == AUTH_WPA2_PSK)
        return ENC_TYPE_CCMP;
    return -1;
}

uint8_t ESP8266WiFiClass::status()
{
    int status = wifi_station_get_connect_status();

    if (status == STATION_GOT_IP)
      return WL_CONNECTED;
    else if (status == STATION_NO_AP_FOUND)
      return WL_NO_SSID_AVAIL;
    else if (status == STATION_CONNECT_FAIL || status == STATION_WRONG_PASSWORD)
      return WL_CONNECT_FAILED;
    else if (status == STATION_IDLE)
      return WL_IDLE_STATUS;
    else
      return WL_DISCONNECTED;
}

void wifi_dns_found_callback(const char *name, ip_addr_t *ipaddr, void *callback_arg)
{
    if (ipaddr)
        (*reinterpret_cast<IPAddress*>(callback_arg)) = ipaddr->addr;
    esp_schedule(); // resume the hostByName function
}

int ESP8266WiFiClass::hostByName(const char* aHostname, IPAddress& aResult)
{
    ip_addr_t addr;
    err_t err = dns_gethostbyname(aHostname, &addr, &wifi_dns_found_callback, &aResult);
    if (err == ERR_OK)
    {
        aResult = addr.addr;
    }
    else if (err == ERR_INPROGRESS)
    {
        esp_yield();
        // will return here when dns_found_callback fires
    }
    else // probably an invalid hostname
    {
        aResult = static_cast<uint32_t>(0);
    }
    return (aResult != 0) ? 1 : 0;
}

ESP8266WiFiClass WiFi;
<commit_msg>WiFi scanning fixes<commit_after>/*
  WiFi.cpp - Library for Arduino Wifi shield.
  Copyright (c) 2011-2014 Arduino.  All right reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.

  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.

  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include "ESP8266WiFi.h"
extern "C" {
#include "c_types.h"
#include "ets_sys.h"
#include "os_type.h"
#include "osapi.h"
#include "mem.h"
#include "user_interface.h"
#include "lwip/opt.h"
#include "lwip/err.h"
#include "lwip/dns.h"
}

extern "C" void esp_schedule();
extern "C" void esp_yield();

ESP8266WiFiClass::ESP8266WiFiClass()
{

}

void ESP8266WiFiClass::mode(WiFiMode m)
{
    ETS_UART_INTR_DISABLE();
    wifi_set_opmode(m);
    ETS_UART_INTR_ENABLE();
}

int ESP8266WiFiClass::begin(const char* ssid)
{
    return begin(ssid, 0);
}

int ESP8266WiFiClass::begin(const char* ssid, const char *passphrase)
{
    if (wifi_get_opmode() == WIFI_AP)
    {
        // turn on AP+STA mode
        mode(WIFI_AP_STA);
    }

    struct station_config conf;
    strcpy(reinterpret_cast<char*>(conf.ssid), ssid);
    if (passphrase)
        strcpy(reinterpret_cast<char*>(conf.password), passphrase);
    else
        *conf.password = 0;

    ETS_UART_INTR_DISABLE();
    wifi_station_set_config(&conf);
    wifi_station_connect();
    ETS_UART_INTR_ENABLE();
    wifi_station_dhcpc_start();
    return 0;
}

void ESP8266WiFiClass::config(IPAddress local_ip, IPAddress gateway, IPAddress subnet)
{
    struct ip_info info;
    info.ip.addr = static_cast<uint32_t>(local_ip);
    info.gw.addr = static_cast<uint32_t>(gateway);
    info.netmask.addr = static_cast<uint32_t>(subnet);

    wifi_station_dhcpc_stop();
    wifi_set_ip_info(STATION_IF, &info);
}

int ESP8266WiFiClass::disconnect()
{
    struct station_config conf;
    *conf.ssid = 0;
    *conf.password = 0;
    ETS_UART_INTR_DISABLE();
    wifi_station_set_config(&conf);
    wifi_station_disconnect();
    ETS_UART_INTR_ENABLE();
    return 0;
}

void ESP8266WiFiClass::softAP(const char* ssid)
{
    softAP(ssid, 0);
}

void ESP8266WiFiClass::softAP(const char* ssid, const char* passphrase)
{
    if (wifi_get_opmode() == WIFI_STA)
    {
        // turn on AP+STA mode
        mode(WIFI_AP_STA);
    }

    struct softap_config conf;
    wifi_softap_get_config(&conf);
    strcpy(reinterpret_cast<char*>(conf.ssid), ssid);
    conf.channel = 1;

    if (!passphrase || strlen(passphrase) == 0)
    {
        conf.authmode = AUTH_OPEN;
        *conf.password = 0;
    }
    else
    {
        conf.authmode = AUTH_WPA2_PSK;
        strcpy(reinterpret_cast<char*>(conf.password), passphrase);
    }

    ETS_UART_INTR_DISABLE();
    wifi_softap_set_config(&conf);
    ETS_UART_INTR_ENABLE();
}

uint8_t* ESP8266WiFiClass::macAddress(uint8_t* mac)
{
    wifi_get_macaddr(STATION_IF, mac);
    return mac;
}

uint8_t* softAPmacAddress(uint8_t* mac)
{
    wifi_get_macaddr(SOFTAP_IF, mac);
    return mac;
}

IPAddress ESP8266WiFiClass::localIP()
{
    struct ip_info ip;
    wifi_get_ip_info(STATION_IF, &ip);
    return IPAddress(ip.ip.addr);
}

IPAddress softAPIP()
{
    struct ip_info ip;
    wifi_get_ip_info(SOFTAP_IF, &ip);
    return IPAddress(ip.ip.addr);
}

IPAddress ESP8266WiFiClass::subnetMask()
{
    struct ip_info ip;
    wifi_get_ip_info(STATION_IF, &ip);
    return IPAddress(ip.netmask.addr);
}

IPAddress ESP8266WiFiClass::gatewayIP()
{
    struct ip_info ip;
    wifi_get_ip_info(STATION_IF, &ip);
    return IPAddress(ip.gw.addr);
}

char* ESP8266WiFiClass::SSID()
{
    static struct station_config conf;
    wifi_station_get_config(&conf);
    return reinterpret_cast<char*>(conf.ssid);
}

// uint8_t* ESP8266WiFiClass::BSSID(uint8_t* bssid)
// {
//  uint8_t* _bssid = WiFiDrv::getCurrentBSSID();
//  memcpy(bssid, _bssid, WL_MAC_ADDR_LENGTH);
//     return bssid;
// }

// int32_t ESP8266WiFiClass::RSSI()
// {
//     return WiFiDrv::getCurrentRSSI();
// }

// uint8_t ESP8266WiFiClass::encryptionType()
// {
//     return WiFiDrv::getCurrentEncryptionType();
// }

void ESP8266WiFiClass::_scanDone(void* result, int status)
{
    if (status != OK)
    {
        ESP8266WiFiClass::_scanCount = 0;
        ESP8266WiFiClass::_scanResult = 0;
    }
    else
    {
      ESP8266WiFiClass::_scanResult = result;
      int i = 0;
      for (bss_info* it = reinterpret_cast<bss_info*>(result); it; it = STAILQ_NEXT(it, next), ++i);
      ESP8266WiFiClass::_scanCount = i;
    }
    esp_schedule();
}

int8_t ESP8266WiFiClass::scanNetworks()
{
    if (wifi_get_opmode() == WIFI_AP)
    {
        mode(WIFI_AP_STA);
    }
    int status = wifi_station_get_connect_status();
    if (status != STATION_GOT_IP && status != STATION_IDLE)
    {
        disconnect();
    }

    if (ESP8266WiFiClass::_scanResult)
    {
        bss_info* it = reinterpret_cast<bss_info*>(ESP8266WiFiClass::_scanResult);
        ESP8266WiFiClass::_scanResult = 0;
        ESP8266WiFiClass::_scanCount = 0;
        while(it)
        {
            bss_info* next = STAILQ_NEXT(it, next);
            os_free(it);
            it = next;
        }
    }

    struct scan_config config;
    config.ssid = 0;
    config.bssid = 0;
    config.channel = 0;
    wifi_station_scan(&config, reinterpret_cast<scan_done_cb_t>(&ESP8266WiFiClass::_scanDone));
    esp_yield();
    return ESP8266WiFiClass::_scanCount;
}

void * ESP8266WiFiClass::_getScanInfoByIndex(int i)
{
    if (!ESP8266WiFiClass::_scanResult || i > ESP8266WiFiClass::_scanCount)
        return 0;

    struct bss_info* it = reinterpret_cast<struct bss_info*>(ESP8266WiFiClass::_scanResult);
    for (; i && it; --i)
        it = STAILQ_NEXT(it, next);

    if (!it)
        return 0;

    return it;
}

char* ESP8266WiFiClass::SSID(uint8_t i)
{
    struct bss_info* it = reinterpret_cast<struct bss_info*>(_getScanInfoByIndex(i));
    if (!it)
        return 0;

    return reinterpret_cast<char*>(it->ssid);
}

int32_t ESP8266WiFiClass::RSSI(uint8_t i)
{
    struct bss_info* it = reinterpret_cast<struct bss_info*>(_getScanInfoByIndex(i));
    if (!it)
        return 0;

    return it->rssi;
}

uint8_t ESP8266WiFiClass::encryptionType(uint8_t i)
{
    struct bss_info* it = reinterpret_cast<struct bss_info*>(_getScanInfoByIndex(i));
    if (!it)
        return -1;

    int authmode = it->authmode;
    if (authmode == AUTH_OPEN)
        return ENC_TYPE_NONE;
    if (authmode == AUTH_WEP)
        return ENC_TYPE_WEP;
    if (authmode == AUTH_WPA_PSK || authmode == AUTH_WPA_WPA2_PSK)  // fixme: is this correct?
        return ENC_TYPE_TKIP;
    if (authmode == AUTH_WPA2_PSK)
        return ENC_TYPE_CCMP;
    return -1;
}

uint8_t ESP8266WiFiClass::status()
{
    int status = wifi_station_get_connect_status();

    if (status == STATION_GOT_IP)
      return WL_CONNECTED;
    else if (status == STATION_NO_AP_FOUND)
      return WL_NO_SSID_AVAIL;
    else if (status == STATION_CONNECT_FAIL || status == STATION_WRONG_PASSWORD)
      return WL_CONNECT_FAILED;
    else if (status == STATION_IDLE)
      return WL_IDLE_STATUS;
    else
      return WL_DISCONNECTED;
}

void wifi_dns_found_callback(const char *name, ip_addr_t *ipaddr, void *callback_arg)
{
    if (ipaddr)
        (*reinterpret_cast<IPAddress*>(callback_arg)) = ipaddr->addr;
    esp_schedule(); // resume the hostByName function
}

int ESP8266WiFiClass::hostByName(const char* aHostname, IPAddress& aResult)
{
    ip_addr_t addr;
    aResult = static_cast<uint32_t>(0);
    err_t err = dns_gethostbyname(aHostname, &addr, &wifi_dns_found_callback, &aResult);
    if (err == ERR_OK)
    {
        aResult = addr.addr;
    }
    else if (err == ERR_INPROGRESS)
    {
        esp_yield();
        // will return here when dns_found_callback fires
    }

    return (aResult != 0) ? 1 : 0;
}

size_t ESP8266WiFiClass::_scanCount = 0;
void* ESP8266WiFiClass::_scanResult = 0;

ESP8266WiFiClass WiFi;
<|endoftext|>


<commit_before>#include "infer.hpp"
#include "ast.hpp"
#include "package.hpp"

#include <sstream>

namespace type {

  // try to open type `self` from signatures in `s`
  static mono open(const ref<state>& s, const mono& self);

  template<class Func>
  static void iter_rows(mono self, Func func) {
    assert(self.kind() == kind::row());
    self.match
      ([&](const app& self) {
        const auto e = extension::unpack(self);
        func(e.attr, e.head);
        iter_rows(e.tail, func);
      }, [](const mono& ) { });
  }

  struct let {
    const ::list<ast::bind> defs;
    const ast::expr body;

    static symbol fix() { return "__fix__"; }

    // rewrite let as non-recursive let + fix **when defining functions**
    static let rewrite(const ast::let& self) {
      using ::list;
      const list<ast::bind> defs = map(self.defs, [](ast::bind self) {
          return self.value.match<ast::bind>
          ([&](const ast::abs& abs) {
            return ast::bind{self.name,
                             ast::app{ast::var{fix()},
                                      ast::abs{self.name >>= list<ast::abs::arg>(),
                                               self.value}
                                      >>= list<ast::expr>()}};
          },
            [&](const ast::expr& expr) { return self; });
        });

      return {defs, *self.body};
    }
  };

  // rewrite app as nested unary applications
  static ast::app rewrite(const ast::app& self) {
    const ast::expr& init = *self.func;
    return foldl(init, self.args, [](ast::expr func, ast::expr arg) {
        return ast::app(func, arg >>= list<ast::expr>());
      }).cast<ast::app>();
  }

  // reconstruct actual type from reified type: ... -> (type 'a) yields 'a
  // note: t must be substituted
  static mono reconstruct(ref<state> s, mono t) {
    // std::clog << "reconstructing: " << s->generalize(t) << std::endl;

    if(auto self = t.get<app>()) {
      // std::clog << "ctor: " << s->generalize((*self)->ctor) << std::endl;

      if((*self)->ctor == ty) {
        return (*self)->arg;
      }
    }

    auto from = s->fresh();
    auto to = s->fresh();

    s->unify(from >>= to, t);
    return reconstruct(s, s->substitute(to));
  }

  // obtain constructor for a monotype
  static cst constructor(mono t) {
    if(auto self = t.get<app>()) {
      return constructor((*self)->ctor);
    }

    if(auto self = t.get<cst>()) {
      return *self;
    }

    // TODO restrictive?
    throw error("constructor must be a constant");
  }

  // try to open type `self` from signatures in `s`
  static mono open(const ref<state>& s, const mono& self) {
    // obtain type constructor from argument type
    const cst c = constructor(s->substitute(self));

    try {
      auto sig = s->sigs->at(c);

      const mono inner = s->fresh();
      s->unify(self >>= inner, s->instantiate(sig));

      return inner;
    } catch(std::out_of_range&) {
      std::stringstream ss;
      ss << "constructor " << tool::quote(c->name.get()) << " has not associated signature";
      throw error(ss.str());
    }
  }

  ////////////////////////////////////////////////////////////////////////////////

  // error: unimplemented
  template<class T>
  static mono infer(const ref<state>& s, const T& self) {
    throw std::logic_error("infer unimplemented: " + tool::type_name(typeid(T)));
  }

  // var
  static mono infer(const ref<state>& s, const ast::var& self) {
    try {
      return s->instantiate(s->vars->find(self.name));
    } catch(std::out_of_range& e) {
      throw error("unbound variable " + tool::quote(self.name.get()));
    }
  }

  // initialize function scope with argument types
  static ref<state> function_scope(const ref<state>& s, const list<ast::abs::arg>& args) {
    auto sub = scope(s);

    for(auto arg: args) {
      const mono type = arg.match<mono>([&](symbol self) {
          return s->fresh();
        },
        [&](ast::abs::typed self) {
          // obtain reified type from annoatation
          const mono reified = infer(s, self.type);

          // TODO do we need gen/inst here?

          // extract concrete type from reified type
          const mono concrete = reconstruct(s, s->substitute(reified));

          return concrete;
        });

      sub->def(arg.name(), type);
    }

    return sub;
  }

  // abs
  static mono infer(const ref<state>& s, const ast::abs& self) {

    // function scope
    const auto sub = function_scope(s, self.args);

    const mono result = s->fresh();
    const mono sig = foldr(result, self.args, [&](auto arg, mono rhs) {
        const mono type = s->instantiate(sub->vars->find(arg.name()));
        return type >>= rhs;
      });

    // infer lambda body with augmented environment
    s->unify(result, infer(sub, *self.body));

    return sig;
  }

  // app
  static mono infer(const ref<state>& s, const ast::app& self) {
    // normalize application as unary
    const ast::app rw = rewrite(self);
    assert(size(rw.args) == 1);

    // infer func/arg types for application
    const auto with_inferred = [&](auto cont) {
      const mono func = infer(s, *rw.func);
      const mono arg = infer(s, rw.args->head);

      return cont(func, arg);
    };

    // check if application works with given func/arg type
    const auto check = [&](mono func, mono arg) {
      const mono ret = s->fresh();
      s->unify(func , arg >>= ret);
      return ret;
    };

    // TODO find a less stupid way of trying all cases?
    try {
      // normal case
      return with_inferred([&](mono func, mono arg) {
          return check(func, arg);
        });

    } catch(error& e) {

      try {
        // open func type and retry
        return with_inferred([&](mono func, mono arg) {
            return check(open(s, func), arg);
          });
      }
      catch(error& ) {  }

      try {
        // open arg type and retry
        return with_inferred([&](mono func, mono arg) {
            return check(func, open(s, arg));
          });
      }
      catch(error& ) {  }

      try {
        // open both func and arg types and retry
        return with_inferred([&](mono func, mono arg) {
            return check(open(s, func), open(s, arg));
          });
      }
      catch(error& ) {  }

      throw e;
    }

  }

  // non-recursive let
  static mono infer(const ref<state>& s, const let& self) {
    auto sub = scope(s);

    for(ast::bind def : self.defs) {
      sub->vars->locals.emplace(def.name, s->generalize(infer(s, def.value)));
    }

    return infer(sub, self.body);
  }

  // recursive let
  static mono infer(const ref<state>& s, const ast::let& self) {
    auto sub = scope(s);

    const mono a = sub->fresh();
    sub->def(let::fix(), (a >>= a) >>= a);

    return infer(sub, let::rewrite(self));
  }

  // sel
  static mono infer(const ref<state>& s, const ast::sel& self) {
    const mono tail = s->fresh(kind::row());
    const mono head = s->fresh(kind::term());

    const mono row = ext(self.name)(head)(tail);
    return record(row) >>= head;
  }

  // record attrs
  static mono infer(const ref<state>& s, const list<ast::record::attr>& attrs) {
    const mono init = empty;
    return foldr(init, attrs, [&](ast::record::attr attr, mono tail) {
        return ext(attr.name)(infer(s, attr.value))(tail);
      });
  }

  // record
  static mono infer(const ref<state>& s, const ast::record& self) {
    return record(infer(s, self.attrs));
  }

  // make
  static mono infer(const ref<state>& s, const ast::make& self) {
    // get signature type
    const poly sig = s->vars->find(self.type);

    auto sub = scope(s);

    const mono outer = s->fresh();
    const mono inner = sub->fresh();

    // instantiate signature at sub level prevents generalization of
    // contravariant side (variables only appearing in the covariant side will
    // be generalized)
    s->unify(ty(outer) >>= ty(inner), sub->instantiate(sig));

    // vanilla covariant type
    const poly reference = sub->generalize(inner);

    // build provided type
    const mono init = empty;
    const mono provided =
      record(foldr(init, self.attrs,
                   [&](const ast::record::attr attr, mono tail) {
                     return row(attr.name, infer(sub, attr.value)) |= tail;
                   }));

    // now also unify inner with provided type
    s->unify(inner, provided);

    // now generalize the provided type
    const poly gen = sub->generalize(inner);

    // generalization check: quantified variables in reference/gen should
    // substitute to the same variables
    std::set<var> quantified;
    for(const var& v : gen.forall) {
      assert(sub->substitute(v) == v);
      quantified.insert(v);
    }

    // make sure all reference quantified references substitute to quantified
    // variables
    for(const var& v : reference.forall) {
      const mono vs = sub->substitute(v);
      if(auto u = vs.get<var>()) {
        auto it = quantified.find(*u);
        if(it != quantified.end()) {
          continue;
        }
      }

      std::stringstream ss;
      logger(ss) << "failed to generalize " << gen
                 << " as " << reference;

      throw error(ss.str());
    }

    return outer;
  }

  // cond
  static mono infer(const ref<state>& s, const ast::cond& self) {
    const mono test = infer(s, *self.test);
    s->unify(test, boolean);

    const mono conseq = infer(s, *self.conseq);
    const mono alt = infer(s, *self.alt);

    const mono result = s->fresh();
    s->unify(result, conseq);
    s->unify(result, alt);

    return result;
  }

  // def
  static mono infer(const ref<state>& s, const ast::def& self) {
    const mono value =
      infer(s, ast::let(ast::bind{self.name, *self.value} >>= list<ast::bind>(),
                        ast::var{self.name}));
    try {
      s->def(self.name, value);
      return io(unit);
    } catch(std::runtime_error& e) {
      throw error(e.what());
    }
  }

  // use
  static mono infer(const ref<state>& s, const ast::use& self) {
    // infer value type
    const mono value = infer(s, *self.env);

    // make sure value type is a record
    const mono row = s->fresh(kind::row());
    s->unify(value, record(row));

    auto sub = scope(s);

    // fill sub scope with record contents
    iter_rows(s->substitute(row), [&](symbol attr, mono t) {
        // TODO generalization issue?
        sub->def(attr, t);
      });

    return infer(sub, *self.body);
  }

  // import
  static mono infer(const ref<state>& s, const ast::import& self) {
    auto it = s->vars->locals.find(self.package);
    if(it != s->vars->locals.end()) {
      throw error("variable " + tool::quote(self.package.get()) + " already defined");
    }

    const package pkg = package::import(self.package);

    s->vars->locals.emplace(self.package, pkg.sig());
    return io(unit);
  }

  // module
  static mono infer(const ref<state>& s, const ast::module& self) {

    // module scope
    const auto sub = function_scope(s, self.args);

    const mono result = s->fresh();
    const mono sig = foldr(result, self.args, [&](auto arg, mono rhs) {
        const mono type = s->instantiate(sub->vars->find(arg.name()));
        return type >>= rhs;
      });

    // infer type for module attributes in module scope
    const mono attrs = infer(sub, self.attrs);

    iter_rows(attrs, [&](symbol attr, mono reified) {
        // make sure each attribute has a reified type
        // const mono type = reconstruct(sub, reified);
      });

    // TODO compute kind from signature

    // TODO create a new type constructor with given name/computed kind

    // in the meantime, return signature
    s->unify(result, record(attrs));
    return sig;
  }

  // lit
  static mono infer(const ref<state>&, const ast::lit<::boolean>& self) {
    return boolean;
  }

  static mono infer(const ref<state>&, const ast::lit<::integer>& self) {
    return integer;
  }

  static mono infer(const ref<state>&, const ast::lit<::real>& self) {
    return real;
  }

  ////////////////////////////////////////////////////////////////////////////////

  struct infer_visitor {
    using type = mono;

    template<class T>
    mono operator()(const T& self, const ref<state>& s) const {
      return infer(s, self);
    }

  };

  mono infer(const ref<state>& s, const ast::expr& self) {
    return self.visit(infer_visitor(), s);
  }

}
<commit_msg>better type reconstruction<commit_after>#include "infer.hpp"
#include "ast.hpp"
#include "package.hpp"

#include <sstream>

namespace type {

  // try to open type `self` from signatures in `s`
  static mono open(const ref<state>& s, const mono& self);

  template<class Func>
  static void iter_rows(mono self, Func func) {
    assert(self.kind() == kind::row());
    self.match
      ([&](const app& self) {
        const auto e = extension::unpack(self);
        func(e.attr, e.head);
        iter_rows(e.tail, func);
      }, [](const mono& ) { });
  }

  struct let {
    const ::list<ast::bind> defs;
    const ast::expr body;

    static symbol fix() { return "__fix__"; }

    // rewrite let as non-recursive let + fix **when defining functions**
    static let rewrite(const ast::let& self) {
      using ::list;
      const list<ast::bind> defs = map(self.defs, [](ast::bind self) {
          return self.value.match<ast::bind>
          ([&](const ast::abs& abs) {
            return ast::bind{self.name,
                             ast::app{ast::var{fix()},
                                      ast::abs{self.name >>= list<ast::abs::arg>(),
                                               self.value}
                                      >>= list<ast::expr>()}};
          },
            [&](const ast::expr& expr) { return self; });
        });

      return {defs, *self.body};
    }
  };

  // rewrite app as nested unary applications
  static ast::app rewrite(const ast::app& self) {
    const ast::expr& init = *self.func;
    return foldl(init, self.args, [](ast::expr func, ast::expr arg) {
        return ast::app(func, arg >>= list<ast::expr>());
      }).cast<ast::app>();
  }

  // reconstruct actual type from reified type: ... -> (type 'a) yields 'a
  // note: t must be substituted
  static mono reconstruct(ref<state> s, mono t) {
    // std::clog << "reconstructing: " << s->generalize(t) << std::endl;
    auto a = s->fresh();

    try {
      s->unify(ty(a), t);
      return a;
      // TODO don't use exceptions for control flow lol
    } catch(error& e) {
      auto from = s->fresh();
      auto to = s->fresh();

      s->unify(from >>= to, t);
      return reconstruct(s, s->substitute(to));
    }

  }

  // obtain constructor for a monotype
  static cst constructor(mono t) {
    if(auto self = t.get<app>()) {
      return constructor((*self)->ctor);
    }

    if(auto self = t.get<cst>()) {
      return *self;
    }

    // TODO restrictive?
    throw error("constructor must be a constant");
  }

  // try to open type `self` from signatures in `s`
  static mono open(const ref<state>& s, const mono& self) {
    // obtain type constructor from argument type
    const cst c = constructor(s->substitute(self));

    try {
      auto sig = s->sigs->at(c);

      const mono inner = s->fresh();
      s->unify(self >>= inner, s->instantiate(sig));

      return inner;
    } catch(std::out_of_range&) {
      std::stringstream ss;
      ss << "constructor " << tool::quote(c->name.get()) << " has not associated signature";
      throw error(ss.str());
    }
  }

  ////////////////////////////////////////////////////////////////////////////////

  // error: unimplemented
  template<class T>
  static mono infer(const ref<state>& s, const T& self) {
    throw std::logic_error("infer unimplemented: " + tool::type_name(typeid(T)));
  }

  // var
  static mono infer(const ref<state>& s, const ast::var& self) {
    try {
      return s->instantiate(s->vars->find(self.name));
    } catch(std::out_of_range& e) {
      throw error("unbound variable " + tool::quote(self.name.get()));
    }
  }

  // initialize function scope with argument types
  static ref<state> function_scope(const ref<state>& s, const list<ast::abs::arg>& args) {
    auto sub = scope(s);

    for(auto arg: args) {
      const mono type = arg.match<mono>([&](symbol self) {
          return s->fresh();
        },
        [&](ast::abs::typed self) {
          // obtain reified type from annoatation
          const mono reified = infer(s, self.type);

          // TODO do we need gen/inst here?

          // extract concrete type from reified type
          const mono concrete = reconstruct(s, s->substitute(reified));

          return concrete;
        });

      sub->def(arg.name(), type);
    }

    return sub;
  }

  // abs
  static mono infer(const ref<state>& s, const ast::abs& self) {

    // function scope
    const auto sub = function_scope(s, self.args);

    const mono result = s->fresh();
    const mono sig = foldr(result, self.args, [&](auto arg, mono rhs) {
        const mono type = s->instantiate(sub->vars->find(arg.name()));
        return type >>= rhs;
      });

    // infer lambda body with augmented environment
    s->unify(result, infer(sub, *self.body));

    return sig;
  }

  // app
  static mono infer(const ref<state>& s, const ast::app& self) {
    // normalize application as unary
    const ast::app rw = rewrite(self);
    assert(size(rw.args) == 1);

    // infer func/arg types for application
    const auto with_inferred = [&](auto cont) {
      const mono func = infer(s, *rw.func);
      const mono arg = infer(s, rw.args->head);

      return cont(func, arg);
    };

    // check if application works with given func/arg type
    const auto check = [&](mono func, mono arg) {
      const mono ret = s->fresh();
      s->unify(func , arg >>= ret);
      return ret;
    };

    // TODO find a less stupid way of trying all cases?
    try {
      // normal case
      return with_inferred([&](mono func, mono arg) {
          return check(func, arg);
        });

    } catch(error& e) {

      try {
        // open func type and retry
        return with_inferred([&](mono func, mono arg) {
            return check(open(s, func), arg);
          });
      }
      catch(error& ) {  }

      try {
        // open arg type and retry
        return with_inferred([&](mono func, mono arg) {
            return check(func, open(s, arg));
          });
      }
      catch(error& ) {  }

      try {
        // open both func and arg types and retry
        return with_inferred([&](mono func, mono arg) {
            return check(open(s, func), open(s, arg));
          });
      }
      catch(error& ) {  }

      throw e;
    }

  }

  // non-recursive let
  static mono infer(const ref<state>& s, const let& self) {
    auto sub = scope(s);

    for(ast::bind def : self.defs) {
      sub->vars->locals.emplace(def.name, s->generalize(infer(s, def.value)));
    }

    return infer(sub, self.body);
  }

  // recursive let
  static mono infer(const ref<state>& s, const ast::let& self) {
    auto sub = scope(s);

    const mono a = sub->fresh();
    sub->def(let::fix(), (a >>= a) >>= a);

    return infer(sub, let::rewrite(self));
  }

  // sel
  static mono infer(const ref<state>& s, const ast::sel& self) {
    const mono tail = s->fresh(kind::row());
    const mono head = s->fresh(kind::term());

    const mono row = ext(self.name)(head)(tail);
    return record(row) >>= head;
  }

  // record attrs
  static mono infer(const ref<state>& s, const list<ast::record::attr>& attrs) {
    const mono init = empty;
    return foldr(init, attrs, [&](ast::record::attr attr, mono tail) {
        return ext(attr.name)(infer(s, attr.value))(tail);
      });
  }

  // record
  static mono infer(const ref<state>& s, const ast::record& self) {
    return record(infer(s, self.attrs));
  }

  // make
  static mono infer(const ref<state>& s, const ast::make& self) {
    // get signature type
    const poly sig = s->vars->find(self.type);

    auto sub = scope(s);

    const mono outer = s->fresh();
    const mono inner = sub->fresh();

    // instantiate signature at sub level prevents generalization of
    // contravariant side (variables only appearing in the covariant side will
    // be generalized)
    s->unify(ty(outer) >>= ty(inner), sub->instantiate(sig));

    // vanilla covariant type
    const poly reference = sub->generalize(inner);

    // build provided type
    const mono init = empty;
    const mono provided =
      record(foldr(init, self.attrs,
                   [&](const ast::record::attr attr, mono tail) {
                     return row(attr.name, infer(sub, attr.value)) |= tail;
                   }));

    // now also unify inner with provided type
    s->unify(inner, provided);

    // now generalize the provided type
    const poly gen = sub->generalize(inner);

    // generalization check: quantified variables in reference/gen should
    // substitute to the same variables
    std::set<var> quantified;
    for(const var& v : gen.forall) {
      assert(sub->substitute(v) == v);
      quantified.insert(v);
    }

    // make sure all reference quantified references substitute to quantified
    // variables
    for(const var& v : reference.forall) {
      const mono vs = sub->substitute(v);
      if(auto u = vs.get<var>()) {
        auto it = quantified.find(*u);
        if(it != quantified.end()) {
          continue;
        }
      }

      std::stringstream ss;
      logger(ss) << "failed to generalize " << gen
                 << " as " << reference;

      throw error(ss.str());
    }

    return outer;
  }

  // cond
  static mono infer(const ref<state>& s, const ast::cond& self) {
    const mono test = infer(s, *self.test);
    s->unify(test, boolean);

    const mono conseq = infer(s, *self.conseq);
    const mono alt = infer(s, *self.alt);

    const mono result = s->fresh();
    s->unify(result, conseq);
    s->unify(result, alt);

    return result;
  }

  // def
  static mono infer(const ref<state>& s, const ast::def& self) {
    const mono value =
      infer(s, ast::let(ast::bind{self.name, *self.value} >>= list<ast::bind>(),
                        ast::var{self.name}));
    try {
      s->def(self.name, value);
      return io(unit);
    } catch(std::runtime_error& e) {
      throw error(e.what());
    }
  }

  // use
  static mono infer(const ref<state>& s, const ast::use& self) {
    // infer value type
    const mono value = infer(s, *self.env);

    // make sure value type is a record
    const mono row = s->fresh(kind::row());
    s->unify(value, record(row));

    auto sub = scope(s);

    // fill sub scope with record contents
    iter_rows(s->substitute(row), [&](symbol attr, mono t) {
        // TODO generalization issue?
        sub->def(attr, t);
      });

    return infer(sub, *self.body);
  }

  // import
  static mono infer(const ref<state>& s, const ast::import& self) {
    auto it = s->vars->locals.find(self.package);
    if(it != s->vars->locals.end()) {
      throw error("variable " + tool::quote(self.package.get()) + " already defined");
    }

    const package pkg = package::import(self.package);

    s->vars->locals.emplace(self.package, pkg.sig());
    return io(unit);
  }

  // module
  static mono infer(const ref<state>& s, const ast::module& self) {

    // module scope
    const auto sub = function_scope(s, self.args);

    const mono result = s->fresh();
    const mono sig = foldr(result, self.args, [&](auto arg, mono rhs) {
        const mono type = s->instantiate(sub->vars->find(arg.name()));
        return type >>= rhs;
      });

    // infer type for module attributes in module scope
    const mono attrs = infer(sub, self.attrs);

    iter_rows(attrs, [&](symbol attr, mono reified) {
        // make sure each attribute has a reified type
        const mono type = reconstruct(sub, reified);
      });

    // TODO compute kind from signature

    // TODO create a new type constructor with given name/computed kind

    // in the meantime, return signature
    s->unify(result, record(attrs));
    return sig;
  }

  // lit
  static mono infer(const ref<state>&, const ast::lit<::boolean>& self) {
    return boolean;
  }

  static mono infer(const ref<state>&, const ast::lit<::integer>& self) {
    return integer;
  }

  static mono infer(const ref<state>&, const ast::lit<::real>& self) {
    return real;
  }

  ////////////////////////////////////////////////////////////////////////////////

  struct infer_visitor {
    using type = mono;

    template<class T>
    mono operator()(const T& self, const ref<state>& s) const {
      return infer(s, self);
    }

  };

  mono infer(const ref<state>& s, const ast::expr& self) {
    return self.visit(infer_visitor(), s);
  }

}
<|endoftext|>